#include <sys/types.h>
#include <sys/prctl.h>
#include <sys/mman.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <stdint.h>
#include <signal.h>
#include <string.h>

#define PAGEMAP_ENTRY 8
#define GET_BIT(X,Y) (X & ((uint64_t)1<<Y)) >> Y
#define GET_PFN(X) X & 0x7FFFFFFFFFFFFF

const int __endian_bit = 1;
#define is_bigendian() ( (*(char*)&__endian_bit) == 0 )
static long pgsz;

/*
 * Set the early kill mode reaction state to MCE error.
 */
static void early_reaction() {
   printf("Setting Early kill... ");
   if (prctl(PR_MCE_KILL, PR_MCE_KILL_SET, PR_MCE_KILL_EARLY, 0, 0) == 0)
      printf("Ok\n");
   else
      printf("Failure !\n");
}

/*
 * Return the physical address associated to a given local virtual address,
 * or -1 in case of an error.
 */
static uint64_t physical_address(uint64_t virt_addr) {
   char path_buf [0x100];
   FILE * f;
   uint64_t read_val, file_offset, pfn = 0;
   unsigned char c_buf[PAGEMAP_ENTRY];
   pid_t my_pid = getpid();
   int status, i;

   sprintf(path_buf, "/proc/%u/pagemap", my_pid);
   
   f = fopen(path_buf, "rb");
   if(!f){
      printf("Error! Cannot open %s\n", path_buf);
      return (uint64_t)-1;
   }
   
   file_offset = virt_addr / (uint64_t)pgsz * PAGEMAP_ENTRY;
   status = fseek(f, (long)file_offset, SEEK_SET);
   if(status){
      perror("Failed to do fseek!");
      fclose(f);
      return (uint64_t)-1;
   }

   for(i=0; i < PAGEMAP_ENTRY; i++){
      int c = getc(f);
      if(c==EOF){
         fclose(f);
         return (uint64_t)-1;
      }
      if(is_bigendian())
           c_buf[i] = (unsigned char)c;
      else
           c_buf[PAGEMAP_ENTRY - i - 1] = (unsigned char)c;
   }
   fclose(f);

   read_val = 0;
   for(i=0; i < PAGEMAP_ENTRY; i++){
      read_val = (read_val << 8) + c_buf[i];
   }

   if(GET_BIT(read_val, 63)) {
      pfn = GET_PFN(read_val);
   } else {
      printf("Page not present !\n");
   }
   if(GET_BIT(read_val, 62))
      printf("Page swapped\n");

   if (pfn == 0)
      return (uint64_t)-1;

   return pfn * (uint64_t)pgsz;
}

/*
 * SIGBUS handler to display the given information.
 */
static void sigbus_action(int signum, siginfo_t *siginfo, void *ctx) {
   printf("Signal %d received: ", signum);
   printf("%s on vaddr: %p\n",
      (siginfo->si_code == 4? "BUS_MCEERR_AR":"BUS_MCEERR_AO"),
      siginfo->si_addr);

  if (siginfo->si_code == 4) { /* BUS_MCEERR_AR */
	fprintf(stderr, "Exit from the signal handler on BUS_MCEERR_AR\n");
	_exit(1);
  }
}

int main(int argc, char ** argv) {
   struct sigaction my_sigaction;
   uint64_t virt_addr = 0, phys_addr;
   void *local_pnt;

   // Need to have the CAP_SYS_ADMIN capability to get PFNs values in pagemap.
   if (getuid() != 0) {
      fprintf(stderr, "Usage: %s needs to run as root\n", argv[0]);
      exit(EXIT_FAILURE);
   }

   // attach our SIGBUS handler.
   memset(&my_sigaction, 0, sizeof(my_sigaction));
   my_sigaction.sa_sigaction = sigbus_action;
   my_sigaction.sa_flags = SA_SIGINFO | SA_NODEFER;
   sigemptyset(&my_sigaction.sa_mask);
   if (sigaction(SIGBUS, &my_sigaction, NULL) == -1) {
      perror("Signal handler attach failed");
      exit(EXIT_FAILURE);
   }

   pgsz = sysconf(_SC_PAGESIZE);
   if (pgsz == -1) {
	   perror("sysconf(_SC_PAGESIZE)");
	   exit(EXIT_FAILURE);
   }
   early_reaction();

   // Allocate a private page.
   local_pnt = mmap(NULL, pgsz, PROT_READ|PROT_WRITE, MAP_ANON|MAP_PRIVATE, -1, 0);
   if (local_pnt == MAP_FAILED) {
      fprintf(stderr, "Memory Allocation failed !\n");
      exit(EXIT_FAILURE);
   }
   virt_addr = (uint64_t)local_pnt;

   // Dirty / map the page.
   sprintf((char *)local_pnt, "My page\n");

   phys_addr = physical_address(virt_addr);
   if (phys_addr == -1) {
      fprintf(stderr, "Virtual address translation 0x%llx failed\n", 
         (unsigned long long)virt_addr);
      exit(EXIT_FAILURE);
   }
   printf("\nData pages at 0x%llx  physically 0x%llx\n",
      (unsigned long long)virt_addr, (unsigned long long)phys_addr);
   fflush(stdout);

   printf("\nPress ENTER to continue\n");
   fgetc(stdin);
   
   // read the string at the beginning of page.
   printf("%s", (char *)local_pnt);

   phys_addr = physical_address(virt_addr);
   if (phys_addr == -1) {
      fprintf(stderr, "Virtual address translation 0x%llx failed\n", 
         (unsigned long long)virt_addr);
   } else {
      printf("\nData pages at 0x%llx  physically 0x%llx\n",
         (unsigned long long)virt_addr, (unsigned long long)phys_addr);
   }

   return 0;
}