[asan Win64] Use multiple trampolines for interception

[wang0109]

Summary

[asan Win64] Use multiple trampolines for interception
Due to large address space in x64, it is not safe to assume
a single trampoline area can be reachable within +/-2G of
every one of the intercepted function address. This patch
provides a way to search and allocate a free page from the
end of each image for trampoline usage. Currently it simply
allocates a new page for every interception.

[wang0109]

This CL is superceded and merged by:

http://reviews.llvm.org/D22111

lib/interception/interception_win.cc

 //===----------------------------------------------------------------------===//
 
 #ifdef _WIN32
 
 #include "interception.h"
 #define WIN32_LEAN_AND_MEAN
 #include <windows.h>
 
+#include <Psapi.h>
+#include "sanitizer_common/sanitizer_common.h"
 namespace __interception {
 
 // FIXME: internal_str* and internal_mem* functions should be moved from the
 // ASan sources into interception/.
 
 static void _memset(void *p, int value, size_t sz) {
   for (size_t i = 0; i < sz; ++i)
     ((char*)p)[i] = (char)value;
 }
 
 static void _memcpy(void *dst, void *src, size_t sz) {
   char *dst_c = (char*)dst,
        *src_c = (char*)src;
   for (size_t i = 0; i < sz; ++i)
     dst_c[i] = src_c[i];
 }
 
 #if SANITIZER_WINDOWS64
-static void WriteIndirectJumpInstruction(char *jmp_from, uptr *indirect_target) {  // NOLINT
+static void WriteIndirectJumpInstruction(char *jmp_from, char *indirect_target) {  // NOLINT
   // jmp [rip + XXYYZZWW] = FF 25 WW ZZ YY XX, where
   // XXYYZZWW is an offset from jmp_from.
   // The displacement is still 32-bit in x64, so indirect_target must be located
   // within +/- 2GB range.
-  int offset = (int)(indirect_target - (uptr *)jmp_from);
+
+  int offset = (int)(indirect_target - jmp_from - 6);
   jmp_from[0] = '\xFF';
   jmp_from[1] = '\x25';
   *(int*)(jmp_from + 2) = offset;
 }
 #else
 static void WriteJumpInstruction(char *jmp_from, char *to) {
   // jmp XXYYZZWW = E9 WW ZZ YY XX, where XXYYZZWW is an offset from jmp_from
   // to the next instruction to the destination.
   ptrdiff_t offset = to - jmp_from - 5;
   *jmp_from = '\xE9';
   *(ptrdiff_t*)(jmp_from + 1) = offset;
 }
 #endif
 
 static void WriteTrampolineJumpInstruction(char *jmp_from, char *to) {
 #if SANITIZER_WINDOWS64
   // Emit an indirect jump through immediately following bytes:
   // jmp_from:
   //   jmp [rip + 6]
   //   .quad to
   // Store the address.
-  uptr *indirect_target = (uptr *)(jmp_from + 6);
-  *indirect_target = (uptr)to;
+  char *indirect_target = jmp_from + 6;
+  *(uptr*)indirect_target = (uptr)to;
   // Write the indirect jump.
   WriteIndirectJumpInstruction(jmp_from, indirect_target);
 #else
   WriteJumpInstruction(jmp_from, to);
 #endif
 }
 
-static void WriteInterceptorJumpInstruction(char *jmp_from, char *to) {
+static void WriteInterceptorJumpInstruction(char *jmp_from, char *to, char *pointer_store_addr) {
 #if SANITIZER_WINDOWS64
   // Emit an indirect jump through immediately following bytes:
   // jmp_from:
-  //   jmp [rip - 8]
+  //   jmp [rip + signed offset]
   //   .quad to
-  // Store the address.
-  uptr *indirect_target = (uptr *)(jmp_from - 8);
-  *indirect_target = (uptr)to;
+  char *indirect_target = pointer_store_addr;
+
+  // Check overflow of +/-2G limitation.
+  ptrdiff_t distance =  indirect_target - jmp_from;
+  if (distance > INT32_MAX || distance < INT32_MIN) {
+    __debugbreak();
+  }
+
+  *(uptr*)indirect_target = (uptr)to;
   // Write the indirect jump.
   WriteIndirectJumpInstruction(jmp_from, indirect_target);
 #else
   WriteJumpInstruction(jmp_from, to);
 #endif
 }
 
-static char *GetMemoryForTrampoline(size_t size) {
+// Allocate a new page near end of image of the intercepted function.
+// This allocates new memeory for each call.
+static char *AllocMemoryForTrampoline(char *intercept_addr, size_t size) {
+#if !SANITIZER_WINDOWS64
   // Trampolines are allocated from a common pool.
   const int POOL_SIZE = 1024;
   static char *pool = NULL;
   static size_t pool_used = 0;
   if (!pool) {
     pool = (char *)VirtualAlloc(NULL, POOL_SIZE, MEM_RESERVE | MEM_COMMIT,
                                 PAGE_EXECUTE_READWRITE);
     // FIXME: Might want to apply PAGE_EXECUTE_READ access after all the
     // interceptors are in place.
     if (!pool)
       return NULL;
     _memset(pool, 0xCC /* int 3 */, POOL_SIZE);
   }
 
   if (pool_used + size > POOL_SIZE)
     return NULL;
 
   char *ret = pool + pool_used;
   pool_used += size;
   return ret;
+#else
+  // Because images may spread across the large 64-bit address space, it is not
+  // safe to assume a single trampoline can be reachable within +/- 2G from all
+  // intercepted function addresses. The current approach search for a free
+  // page from the end of each intercepted image, and allocate memory for
+  // trampoline there.
+  static SIZE_T alloc_granularity = 0;
+  static SIZE_T page_size = 0;
+  if (!alloc_granularity) alloc_granularity = __sanitizer::GetMmapGranularity();
+  if (!page_size) page_size = __sanitizer::GetPageSizeCached();
+  if (!alloc_granularity || !page_size) return nullptr;
+
+  // To be returned by this function.
+  char *trampoline_page_begin = 0;
+  // One page is enough for store the head and pointer.
+  // Only alloc one page to save on commit charge.
+  uptr trampoline_alloc_size = page_size;
+  // Give up search for memory page at certain point.
+  static const uptr kRangeLimit = (100ULL) << 20;  // 100 Meg.
+
+  // Identify current image, based on address of intercepted function.
+  HMODULE curr_module = NULL;
+  GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
+                    (LPCTSTR)intercept_addr, &curr_module);
+  if (!curr_module) return nullptr;
+
+  // Get image base address and size.
+  MODULEINFO module_info;
+  memset(&module_info, 0, sizeof(module_info));
+  if (!GetModuleInformation(GetCurrentProcess(), curr_module, &module_info,
+                            sizeof(module_info)))
+    return nullptr;
+
+  // Search starting from end of the image.
+  uptr search_begin =
+      (uptr)module_info.lpBaseOfDll + (uptr)module_info.SizeOfImage;
+
+  for (uptr region_begin = search_begin;
+       region_begin < search_begin + kRangeLimit;) {
+    MEMORY_BASIC_INFORMATION info;
+    if (0 == ::VirtualQuery((LPCVOID)region_begin, &info, sizeof(info)))
+      return nullptr;
+
+    if (info.State == MEM_FREE) {
+      uptr free_begin = (uptr)info.BaseAddress;
+      uptr free_end = free_begin + (uptr)info.RegionSize;
+      // Free region can begin at address non-aligned with granularity,
+      // VirtualAlloc would fail on such address, thus, so round up to the
+      // next granularity. If it is already aligned, round up to itself.
+      uptr next_alloc_begin =
+          ((free_begin - 1) & ~(alloc_granularity - 1)) + alloc_granularity;
+      uptr next_alloc_end =
+          next_alloc_begin +
+          alloc_granularity;  // Min alloc size is one granularity.
+
+      // Enough space?
+      if (free_end >= next_alloc_end) {
+        LPVOID alloced =
+            ::VirtualAlloc((LPVOID)next_alloc_begin, trampoline_alloc_size,
+                           MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
+
+        if (!alloced) return nullptr;
+
+        // Found it.
+        trampoline_page_begin = (char *)alloced;
+        return trampoline_page_begin;
+      }
+    }
+
+    // Proceed to next region.
+    region_begin += info.RegionSize;
+  }
+
+  // Search limit reached and found nothing.
+  return nullptr;
+#endif
 }
 
 // Returns 0 on error.
 static size_t RoundUpToInstrBoundary(size_t size, char *code) {
 #if SANITIZER_WINDOWS64
   // Win64 RoundUpToInstrBoundary is a work in progress.
   size_t cursor = 0;
   while (cursor < size) {
size_t kExtraPrevBytes = 8;  // The extra bytes we need to mark READWRITE for
                                // page access, that is preceeding the begin
                                // of function.
 #else
   size_t kHeadMin = 5;
   size_t kTrampolineJumpSize = 5;
   size_t kExtraPrevBytes = 0;
 #endif
   size_t head = kHeadMin;
+
+  char *trampoline = 0;
   if (orig_old_func) {
     // Find out the number of bytes of the instructions we need to copy
     // to the trampoline and store it in 'head'.
     head = RoundUpToInstrBoundary(kHeadMin, old_bytes);
     if (!head)
       return false;
 
     // Put the needed instructions into the trampoline bytes.
-    char *trampoline = GetMemoryForTrampoline(head + kTrampolineJumpSize);
+    trampoline = AllocMemoryForTrampoline(old_bytes, head + kTrampolineJumpSize);
     if (!trampoline)
       return false;
     _memcpy(trampoline, old_bytes, head);
     WriteTrampolineJumpInstruction(trampoline + head, old_bytes + head);
     *orig_old_func = (uptr)trampoline;
   }
 
   // Now put the "jmp <new_func>" instruction at the original code location.
   // We should preserve the EXECUTE flag as some of our own code might be
   // located in the same page (sic!).  FIXME: might consider putting the
   // __interception code into a separate section or something?
   DWORD old_prot, unused_prot;
   // TODO(wwchrome): Properly handle access violations when finding a safe
   // region to store the indirect jump target address.
   // Need to mark extra 8 bytes for Win64 because jmp [rip -8]
   if (!VirtualProtect((void *)(old_bytes - kExtraPrevBytes),
                       head + kExtraPrevBytes, PAGE_EXECUTE_READWRITE,
                       &old_prot))
     return false;
 
-  WriteInterceptorJumpInstruction(old_bytes, (char *)new_func);
+  if (trampoline == 0)
+    __debugbreak();
+  // By design..  trampoline + 256 bytes, is where pointer is located.
+  static const ptrdiff_t kTrampolineHeadSpace = 256;
+  char *pointer_stored_addr = trampoline + kTrampolineHeadSpace;
+  WriteInterceptorJumpInstruction(old_bytes, (char *)new_func,  pointer_stored_addr);
   _memset(old_bytes + kHeadMin, 0xCC /* int 3 */, head - kHeadMin);
 
   // Restore the original permissions.
   if (!VirtualProtect((void *)(old_bytes - kExtraPrevBytes),
                       head + kExtraPrevBytes, old_prot, &unused_prot))
     return false;  // not clear if this failure bothers us.
 
   return true;