[ASan/Win] Refactor __interception::OverrideFunction a bit, allow the "orig_old_func" parameter to be zero.

[timurrrr]

Reverted the code style changes

[timurrrr]

Reviewed with Alexander offline.

[timurrrr]

r208989.

interception/interception_win.cc

   *(ptrdiff_t*)(jmp_from + 1) = offset;
 }
 
-bool OverrideFunction(uptr old_func, uptr new_func, uptr *orig_old_func) {
-#ifdef _WIN64
-# error OverrideFunction was not tested on x64
-#endif
-  // Basic idea:
-  // We write 5 bytes (jmp-to-new_func) at the beginning of the 'old_func'
-  // to override it. We want to be able to execute the original 'old_func' from
-  // the wrapper, so we need to keep the leading 5+ bytes ('head') of the
-  // original instructions somewhere with a "jmp old_func+head".
-  // We call these 'head'+5 bytes of instructions a "trampoline".
-
+static char *GetMemoryForTrampoline(size_t size) {
   // Trampolines are allocated from a common pool.
   const int POOL_SIZE = 1024;
   static char *pool = NULL;
   static size_t pool_used = 0;
-  if (pool == NULL) {
-    pool = (char*)VirtualAlloc(NULL, POOL_SIZE,
-                               MEM_RESERVE | MEM_COMMIT,
-                               PAGE_EXECUTE_READWRITE);
-    // FIXME: set PAGE_EXECUTE_READ access after setting all interceptors?
-    if (pool == NULL)
-      return false;
+  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);
   }
 
-  char* old_bytes = (char*)old_func;
-  char* trampoline = pool + pool_used;
+  if (pool_used + size > POOL_SIZE)
+    return NULL;
 
-  // Find out the number of bytes of the instructions we need to copy to the
-  // island and store it in 'head'.
-  size_t head = 0;
-  while (head < 5) {
-    switch (old_bytes[head]) {
+  char *ret = pool + pool_used;
+  pool_used += size;
+  return ret;
+}
+
+// Returns 0 on error.
+static size_t RoundUpToInstrBoundary(size_t size, char *code) {
+  size_t cursor = 0;
+  while (cursor < size) {
+    switch (code[cursor]) {
       case '\x51':  // push ecx
       case '\x52':  // push edx
       case '\x53':  // push ebx

       case '\x56':  // push esi
       case '\x57':  // push edi
       case '\x5D':  // pop ebp
-        head++;
+        cursor++;
         continue;
       case '\x6A':  // 6A XX = push XX
-        head += 2;
+        cursor += 2;
         continue;
       case '\xE9':  // E9 XX YY ZZ WW = jmp WWZZYYXX
-        head += 5;
+        cursor += 5;
         continue;
     }
-    switch (*(unsigned short*)(old_bytes + head)) {  // NOLINT
+    switch (*(unsigned short*)(code + cursor)) {  // NOLINT
       case 0xFF8B:  // 8B FF = mov edi, edi
       case 0xEC8B:  // 8B EC = mov ebp, esp
       case 0xC033:  // 33 C0 = xor eax, eax
-        head += 2;
+        cursor += 2;
         continue;
       case 0x458B:  // 8B 45 XX = mov eax, dword ptr [ebp+XXh]
       case 0x5D8B:  // 8B 5D XX = mov ebx, dword ptr [ebp+XXh]
       case 0xEC83:  // 83 EC XX = sub esp, XX
-        head += 3;
+        cursor += 3;
         continue;
       case 0xC1F7:  // F7 C1 XX YY ZZ WW = test ecx, WWZZYYXX
-        head += 6;
+        cursor += 6;
         continue;
       case 0x3D83:  // 83 3D XX YY ZZ WW TT = cmp TT, WWZZYYXX
-        head += 7;
+        cursor += 7;
         continue;
     }
-    switch (0x00FFFFFF & *(unsigned int*)(old_bytes + head)) {
+    switch (0x00FFFFFF & *(unsigned int*)(code + cursor)) {
       case 0x24448A:  // 8A 44 24 XX = mov eal, dword ptr [esp+XXh]
       case 0x244C8B:  // 8B 4C 24 XX = mov ecx, dword ptr [esp+XXh]
       case 0x24548B:  // 8B 54 24 XX = mov edx, dword ptr [esp+XXh]
       case 0x24748B:  // 8B 74 24 XX = mov esi, dword ptr [esp+XXh]
       case 0x247C8B:  // 8B 7C 24 XX = mov edi, dword ptr [esp+XXh]
-        head += 4;
+        cursor += 4;
         continue;
     }
 

     // in visible and readable error messages. However, merely calling abort()
     // or __debugbreak() leads to an infinite recursion in CheckFailed.
     // Do we have a good way to abort with an error message here?
-    return false;
+    return 0;
   }
 
-  if (pool_used + head + 5 > POOL_SIZE)
-    return false;
+  return cursor;
+}
 
-  // Now put the "jump to trampoline" instruction into the original code.
+bool OverrideFunction(uptr old_func, uptr new_func, uptr *orig_old_func) {
+#ifdef _WIN64
+#error OverrideFunction is not yet supported on x64
+#endif
+  // Function overriding works basically like this:
+  // We write "jmp <new_func>" (5 bytes) at the beginning of the 'old_func'
+  // to override it.
+  // We might want to be able to execute the original 'old_func' from the
+  // wrapper, in this case we need to keep the leading 5+ bytes ('head')
+  // of the original code somewhere with a "jmp <old_func+head>".
+  // We call these 'head'+5 bytes of instructions a "trampoline".
+  char *old_bytes = (char *)old_func;
+
+  // We'll need at least 5 bytes for a 'jmp'.
+  size_t head = 5;
+  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(head, old_bytes);
+    if (!head)
+      return false;
+
+    // Put the needed instructions into the trampoline bytes.
+    char *trampoline = GetMemoryForTrampoline(head + 5);
+    if (!trampoline)
+      return false;
+    _memcpy(trampoline, old_bytes, head);
+    WriteJumpInstruction(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;
-  if (!VirtualProtect((void*)old_func, head, PAGE_EXECUTE_READWRITE,
+  if (!VirtualProtect((void *)old_bytes, head, PAGE_EXECUTE_READWRITE,
                       &old_prot))
     return false;
 
-  // Put the needed instructions into the trampoline bytes.
-  _memcpy(trampoline, old_bytes, head);
-  WriteJumpInstruction(trampoline + head, old_bytes + head);
-  *orig_old_func = (uptr)trampoline;
-  pool_used += head + 5;
-
-  // Intercept the 'old_func'.
-  WriteJumpInstruction(old_bytes, (char*)new_func);
+  WriteJumpInstruction(old_bytes, (char *)new_func);
   _memset(old_bytes + 5, 0xCC /* int 3 */, head - 5);
 
-  if (!VirtualProtect((void*)old_func, head, old_prot, &unused_prot))
+  // Restore the original permissions.
+  if (!VirtualProtect((void *)old_bytes, head, old_prot, &unused_prot))
     return false;  // not clear if this failure bothers us.
 
   return true;