tools/memory_watcher/preamble_patcher.cc

   1 // Copyright (c) 2010 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "preamble_patcher.h"
   6 #include "memory_hook.h"
   7 #include "mini_disassembler.h"
   8
   9 // compatibility shims
  10 #include "base/logging.h"
  11
  12 // Definitions of assembly statements we need
  13 #define ASM_JMP32REL 0xE9
  14 #define ASM_INT3 0xCC
  15
  16 namespace sidestep {
  17
  18 SideStepError PreamblePatcher::RawPatchWithStubAndProtections(
  19     void* target_function, void *replacement_function,
  20     unsigned char* preamble_stub, unsigned long stub_size,
  21     unsigned long* bytes_needed) {
  22   // We need to be able to write to a process-local copy of the first
  23   // MAX_PREAMBLE_STUB_SIZE bytes of target_function. We may be giving execute
  24   // privilege to something that doesn't have it, but that's the price to pay
  25   // for tools.
  26   DWORD old_target_function_protect = 0;
  27   BOOL succeeded = ::VirtualProtect(reinterpret_cast<void*>(target_function),
  28                                     MAX_PREAMBLE_STUB_SIZE,
  29                                     PAGE_EXECUTE_READWRITE,
  30                                     &old_target_function_protect);
  31   if (!succeeded) {
  32     ASSERT(false, "Failed to make page containing target function "
  33                    "copy-on-write.");
  34     return SIDESTEP_ACCESS_DENIED;
  35   }
  36
  37   SideStepError error_code = RawPatchWithStub(target_function,
  38                                               replacement_function,
  39                                               preamble_stub,
  40                                               stub_size,
  41                                               bytes_needed);
  42   if (SIDESTEP_SUCCESS != error_code) {
  43     ASSERT1(false);
  44     return error_code;
  45   }
  46
  47   // Restore the protection of the first MAX_PREAMBLE_STUB_SIZE bytes of
  48   // pTargetFunction to what they were before we started goofing around.
  49   succeeded = ::VirtualProtect(reinterpret_cast<void*>(target_function),
  50                                MAX_PREAMBLE_STUB_SIZE,
  51                                old_target_function_protect,
  52                                &old_target_function_protect);
  53   if (!succeeded) {
  54     ASSERT(false, "Failed to restore protection to target function.");
  55     // We must not return an error here because the function has actually
  56     // been patched, and returning an error would likely cause our client
  57     // code not to unpatch it.  So we just keep going.
  58   }
  59
  60   // Flush the instruction cache to make sure the processor doesn't execute the
  61   // old version of the instructions (before our patch).
  62   //
  63   // FlushInstructionCache is actually a no-op at least on single-processor
  64   // XP machines.  I'm not sure why this is so, but it is, yet I want to keep
  65   // the call to the API here for correctness in case there is a difference in
  66   // some variants of Windows/hardware.
  67   succeeded = ::FlushInstructionCache(::GetCurrentProcess(),
  68                                       target_function,
  69                                       MAX_PREAMBLE_STUB_SIZE);
  70   if (!succeeded) {
  71     ASSERT(false, "Failed to flush instruction cache.");
  72     // We must not return an error here because the function has actually
  73     // been patched, and returning an error would likely cause our client
  74     // code not to unpatch it.  So we just keep going.
  75   }
  76
  77   return SIDESTEP_SUCCESS;
  78 }
  79
  80 SideStepError PreamblePatcher::RawPatch(void* target_function,
  81                                         void* replacement_function,
  82                                         void** original_function_stub) {
  83   if (!target_function || !replacement_function || !original_function_stub ||
  84       (*original_function_stub) || target_function == replacement_function) {
  85     ASSERT(false, "Preconditions not met");
  86     return SIDESTEP_INVALID_PARAMETER;
  87   }
  88
  89   // @see MAX_PREAMBLE_STUB_SIZE for an explanation of how we arrives at
  90   // this size
  91   unsigned char* preamble_stub =
  92     reinterpret_cast<unsigned char*>(
  93       MemoryHook::Alloc(sizeof(unsigned char) * MAX_PREAMBLE_STUB_SIZE));
  94   if (!preamble_stub) {
  95     ASSERT(false, "Unable to allocate preamble-stub.");
  96     return SIDESTEP_INSUFFICIENT_BUFFER;
  97   }
  98
  99   // Change the protection of the newly allocated preamble stub to
 100   // PAGE_EXECUTE_READWRITE. This is required to work with DEP (Data
 101   // Execution Prevention) which will cause an exception if code is executed
 102   // from a page on which you do not have read access.
 103   DWORD old_stub_protect = 0;
 104   BOOL succeeded = VirtualProtect(preamble_stub, MAX_PREAMBLE_STUB_SIZE,
 105                              PAGE_EXECUTE_READWRITE, &old_stub_protect);
 106   if (!succeeded) {
 107     ASSERT(false, "Failed to make page preamble stub read-write-execute.");
 108     delete[] preamble_stub;
 109     return SIDESTEP_ACCESS_DENIED;
 110   }
 111
 112   SideStepError error_code = RawPatchWithStubAndProtections(target_function,
 113                                               replacement_function,
 114                                               preamble_stub,
 115                                               MAX_PREAMBLE_STUB_SIZE,
 116                                               NULL);
 117   if (SIDESTEP_SUCCESS != error_code) {
 118     ASSERT1(false);
 119     delete[] preamble_stub;
 120     return error_code;
 121   }
 122
 123   *original_function_stub = reinterpret_cast<void*>(preamble_stub);
 124
 125   // NOTE: For hooking malloc/free, we don't want to use streams which
 126   // allocate.  Basically, we've hooked malloc, but not necessarily
 127   // hooked free yet.  To do anything which uses the heap could crash
 128   // with a mismatched malloc/free!
 129   //VLOG(1) << "PreamblePatcher::RawPatch successfully patched 0x"
 130   //        << target_function;
 131
 132   return SIDESTEP_SUCCESS;
 133 }
 134
 135 SideStepError PreamblePatcher::Unpatch(void* target_function,
 136                                        void* replacement_function,
 137                                        void* original_function_stub) {
 138   ASSERT1(target_function && original_function_stub);
 139   if (!target_function || !original_function_stub) {
 140     return SIDESTEP_INVALID_PARAMETER;
 141   }
 142
 143   // We disassemble the preamble of the _stub_ to see how many bytes we
 144   // originally copied to the stub.
 145   MiniDisassembler disassembler;
 146   unsigned int preamble_bytes = 0;
 147   while (preamble_bytes < 5) {
 148     InstructionType instruction_type = disassembler.Disassemble(
 149         reinterpret_cast<unsigned char*>(original_function_stub) +
 150         preamble_bytes, preamble_bytes);
 151     if (IT_GENERIC != instruction_type) {
 152       ASSERT(false, "Should only have generic instructions in stub!!");
 153       return SIDESTEP_UNSUPPORTED_INSTRUCTION;
 154     }
 155   }
 156
 157   // Before unpatching, target_function should be a JMP to
 158   // replacement_function.  If it's not, then either it's an error, or
 159   // we're falling into the case where the original instruction was a
 160   // JMP, and we patched the jumped_to address rather than the JMP
 161   // itself.  (For instance, if malloc() is just a JMP to __malloc(),
 162   // we patched __malloc() and not malloc().)
 163   unsigned char* target = reinterpret_cast<unsigned char*>(target_function);
 164   while (1) {    // we stop when target is a JMP to replacement_function
 165     if (target[0] != ASM_JMP32REL) {
 166       ASSERT(false, "target_function does not look like it was patched.");
 167       return SIDESTEP_INVALID_PARAMETER;
 168     }
 169     int relative_offset;   // Windows guarantees int is 4 bytes
 170     ASSERT1(sizeof(relative_offset) == 4);
 171     memcpy(reinterpret_cast<void*>(&relative_offset),
 172            reinterpret_cast<void*>(target + 1), 4);
 173     unsigned char* jump_to = target + 5 + relative_offset;
 174     if (jump_to == replacement_function)
 175       break;
 176     target = jump_to;      // follow the jmp
 177   }
 178
 179   // We need to be able to write to a process-local copy of the first
 180   // MAX_PREAMBLE_STUB_SIZE bytes of target_function. We may be giving execute
 181   // privilege to something that doesn't have it, but that's the price to pay
 182   // for tools.
 183   DWORD old_target_function_protect = 0;
 184   BOOL succeeded = ::VirtualProtect(reinterpret_cast<void*>(target),
 185                                     MAX_PREAMBLE_STUB_SIZE,
 186                                     PAGE_EXECUTE_READWRITE,
 187                                     &old_target_function_protect);
 188   if (!succeeded) {
 189     ASSERT(false, "Failed to make page containing target function "
 190                    "copy-on-write.");
 191     return SIDESTEP_ACCESS_DENIED;
 192   }
 193
 194   // Replace the first few bytes of the original function with the bytes we
 195   // previously moved to the preamble stub.
 196   memcpy(reinterpret_cast<void*>(target),
 197          original_function_stub, preamble_bytes);
 198
 199   // Stub is now useless so delete it.
 200   // [csilvers: Commented out for perftools because it causes big problems
 201   //  when we're unpatching malloc.  We just let this live on as a leak.]
 202   //delete original_function_stub;
 203
 204   // Restore the protection of the first MAX_PREAMBLE_STUB_SIZE bytes of
 205   // target to what they were before we started goofing around.
 206   succeeded = ::VirtualProtect(reinterpret_cast<void*>(target),
 207                                MAX_PREAMBLE_STUB_SIZE,
 208                                old_target_function_protect,
 209                                &old_target_function_protect);
 210
 211   // Flush the instruction cache to make sure the processor doesn't execute the
 212   // old version of the instructions (before our patch).
 213   //
 214   // See comment on FlushInstructionCache elsewhere in this file.
 215   succeeded = ::FlushInstructionCache(::GetCurrentProcess(),
 216                                       target,
 217                                       MAX_PREAMBLE_STUB_SIZE);
 218   if (!succeeded) {
 219     ASSERT(false, "Failed to flush instruction cache.");
 220     return SIDESTEP_UNEXPECTED;
 221   }
 222
 223   VLOG(1) << "PreamblePatcher::Unpatch successfully unpatched 0x"
 224           << target_function;
 225   return SIDESTEP_SUCCESS;
 226 }
 227
 228 };  // namespace sidestep