chrome_frame/function_stub.h

   1 // Copyright (c) 2009 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
   6 #ifndef CHROME_FRAME_FUNCTION_STUB_H_
   7 #define CHROME_FRAME_FUNCTION_STUB_H_
   8
   9 #include <windows.h>
  10 #include "base/logging.h"
  11
  12 // IMPORTANT: The struct below must be byte aligned.
  13 #pragma pack(push)
  14 #pragma pack(1)
  15
  16 #ifndef _M_IX86
  17 #error Only x86 supported right now.
  18 #endif
  19
  20 extern "C" IMAGE_DOS_HEADER __ImageBase;
  21
  22 // This struct is assembly code + signature.  The purpose of the struct is to be
  23 // able to hook an existing function with our own and store information such
  24 // as the original function pointer with the code stub.  Typically this is used
  25 // for patching entries of a vtable or e.g. a globally registered wndproc
  26 // for a class as opposed to a window.
  27 // When unhooking, you can just call the BypassStub() function and leave the
  28 // stub in memory.  This unhooks your function while leaving the (potential)
  29 // chain of patches intact.
  30 //
  31 // @note: This class is meant for __stdcall calling convention and
  32 //        it uses eax as a temporary variable.  The struct can
  33 //        be improved in the future to save eax before the
  34 //        operation and then restore it.
  35 //
  36 // For instance if the function prototype is:
  37 //
  38 // @code
  39 //   LRESULT WndProc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam);
  40 // @endcode
  41 //
  42 // and we would like to add one static argument to make it, say:
  43 //
  44 // @code
  45 //   LRESULT MyNewWndProc(WNDPROC original, HWND hwnd, UINT msg,
  46 //                        WPARAM wparam, LPARAM lparam);
  47 // @endcode
  48 //
  49 // That can be achieved by wrapping the function up with a FunctionStub:
  50 //
  51 // @code
  52 //   FunctionStub* stub = FunctionStub::Create(original_wndproc, MyNewWndProc);
  53 //   SetClassLongPtr(wnd, GCLP_WNDPROC, stub->code());
  54 // @endcode
  55 struct FunctionStub {
  56  private:
  57   typedef enum AsmConstants {
  58     POP_EAX = 0x58,
  59     PUSH = 0x68,
  60     PUSH_EAX = 0x50,
  61     JUMP_RELATIVE = 0xE9
  62   };
  63
  64   FunctionStub(uintptr_t extra_argument, void* dest)
  65       : signature_(reinterpret_cast<HMODULE>(&__ImageBase)) {
  66     Opcodes::Hook& hook = code_.hook_;
  67     hook.pop_return_addr_ = POP_EAX;
  68     hook.push_ = PUSH;
  69     hook.arg_ = extra_argument;
  70     hook.push_return_addr_ = PUSH_EAX;
  71     hook.jump_to_ = JUMP_RELATIVE;
  72
  73     // Calculate the target jump to the destination function.
  74     hook.target_ = CalculateRelativeJump(dest, &hook.jump_to_);
  75
  76     // Allow the process to execute this struct as code.
  77     DWORD old_protect = 0;
  78     // Allow reads too since we want read-only member variable access at
  79     // all times.
  80     ::VirtualProtect(this, sizeof(FunctionStub), PAGE_EXECUTE_READ,
  81                      &old_protect);
  82   }
  83
  84   ~FunctionStub() {
  85     // No more execution rights.
  86     DWORD old_protect = 0;
  87     ::VirtualProtect(this, sizeof(FunctionStub), PAGE_READWRITE, &old_protect);
  88   }
  89
  90   // Calculates the target value for a relative jump.
  91   // The function assumes that the size of the opcode is 1 byte.
  92   inline uintptr_t CalculateRelativeJump(const void* absolute_to,
  93                                          const void* absolute_from) const {
  94     return (reinterpret_cast<uintptr_t>(absolute_to) -
  95             reinterpret_cast<uintptr_t>(absolute_from)) -
  96            (sizeof(uint8) + sizeof(uintptr_t));
  97   }
  98
  99   // Does the opposite of what CalculateRelativeJump does, which is to
 100   // calculate back the absolute address that previously was relative to
 101   // some other address.
 102   inline uintptr_t CalculateAbsoluteAddress(const void* relative_to,
 103                                             uintptr_t relative_address) const {
 104     return relative_address + sizeof(uint8) + sizeof(uintptr_t) +
 105            reinterpret_cast<uintptr_t>(relative_to);
 106   }
 107
 108   // Used to identify function stubs that belong to this module.
 109   HMODULE signature_;
 110
 111   // IMPORTANT: Do not change the order of member variables
 112   union Opcodes {
 113     // Use this struct when the stub forwards the call to our hook function
 114     // providing an extra argument.
 115     struct Hook {
 116       uint8 pop_return_addr_;  // pop eax
 117       uint8 push_;             // push arg  ; push...
 118       uintptr_t arg_;          //           ; extra argument
 119       uint8 push_return_addr_; // push eax  ; push the return address
 120       uint8 jump_to_;          // jmp       ; jump...
 121       uintptr_t target_;       //           ; to the hook function
 122     } hook_;
 123     // When the stub is bypassed, we jump directly to a given target,
 124     // usually the originally hooked function.
 125     struct Bypassed {
 126       uint8 jump_to_;          // jmp to
 127       uintptr_t target_;       // relative target.
 128     } bypassed_;
 129   };
 130
 131   Opcodes code_;
 132
 133  public:
 134   // Neutralizes this stub and converts it to a direct jump to a new target.
 135   void BypassStub(void* new_target) {
 136     DWORD old_protect = 0;
 137     // Temporarily allow us to write to member variables
 138     ::VirtualProtect(this, sizeof(FunctionStub), PAGE_EXECUTE_READWRITE,
 139                      &old_protect);
 140
 141     // Now, just change the first 5 bytes to jump directly to the new target.
 142     Opcodes::Bypassed& bypassed = code_.bypassed_;
 143     bypassed.jump_to_ = JUMP_RELATIVE;
 144     bypassed.target_ = CalculateRelativeJump(new_target, &bypassed.jump_to_);
 145
 146     // Restore the previous protection flags.
 147     ::VirtualProtect(this, sizeof(FunctionStub), old_protect, &old_protect);
 148
 149     // Flush the instruction cache just in case.
 150     ::FlushInstructionCache(::GetCurrentProcess(), this, sizeof(FunctionStub));
 151   }
 152
 153   // @returns the argument supplied in the call to @ref Create
 154   inline uintptr_t argument() const {
 155     DCHECK(code_.hook_.pop_return_addr_ == POP_EAX) << "stub has been disabled";
 156     return code_.hook_.arg_;
 157   }
 158
 159   inline bool is_bypassed() const {
 160     return code_.bypassed_.jump_to_ == JUMP_RELATIVE;
 161   }
 162
 163   inline uintptr_t absolute_target() const {
 164     DCHECK(code_.hook_.pop_return_addr_ == POP_EAX) << "stub has been disabled";
 165     return CalculateAbsoluteAddress(
 166         reinterpret_cast<const void*>(&code_.hook_.jump_to_),
 167         code_.hook_.target_);
 168   }
 169
 170   // Returns true if the stub is valid and enabled.
 171   // Don't call this method after bypassing the stub.
 172   inline bool is_valid() const {
 173     return signature_ == reinterpret_cast<HMODULE>(&__ImageBase) &&
 174         code_.hook_.pop_return_addr_ == POP_EAX;
 175   }
 176
 177   inline PROC code() const {
 178     return reinterpret_cast<PROC>(const_cast<Opcodes*>(&code_));
 179   }
 180
 181   // Use to create a new function stub as shown above.
 182   //
 183   // @param extra_argument The static argument to pass to the function.
 184   // @param dest Target function to which the stub applies.
 185   // @returns NULL if an error occurs, otherwise a pointer to the
 186   // function stub.
 187   //
 188   // TODO(tommi): Change this so that VirtualAlloc isn't called for
 189   // every stub.  Instead we should re-use each allocation for
 190   // multiple stubs.  In practice I'm guessing that there would
 191   // only be one allocation per process, since each allocation actually
 192   // allocates at least one page of memory (4K).  Size of FunctionStub
 193   // is 12 bytes, so one page could house 341 function stubs.
 194   // When stubs are created frequently, VirtualAlloc could fail
 195   // and last error is ERROR_NOT_ENOUGH_MEMORY (8).
 196   static FunctionStub* Create(uintptr_t extra_argument, void* dest) {
 197     DCHECK(dest);
 198
 199     // Use VirtualAlloc to get memory for the stub.  This gives us a
 200     // whole page that we don't share with anyone else.
 201     // Initially the memory must be READWRITE to allow for construction
 202     // PAGE_EXECUTE is set in constructor.
 203     FunctionStub* ret = reinterpret_cast<FunctionStub*>(VirtualAlloc(NULL,
 204         sizeof(FunctionStub), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
 205
 206     if (!ret) {
 207       NOTREACHED();
 208     } else {
 209       // Construct
 210       ret->FunctionStub::FunctionStub(extra_argument, dest);
 211     }
 212
 213     return ret;
 214   }
 215
 216   static FunctionStub* FromCode(void* address) {
 217     Opcodes* code = reinterpret_cast<Opcodes*>(address);
 218     if (code->hook_.pop_return_addr_ == POP_EAX) {
 219       FunctionStub* stub = reinterpret_cast<FunctionStub*>(
 220           reinterpret_cast<uint8*>(address) - sizeof(HMODULE));
 221       if (stub->is_valid())
 222         return stub;
 223     }
 224
 225     return NULL;
 226   }
 227
 228   // Deallocates a FunctionStub.  The stub must not be in use on any thread!
 229   static bool Destroy(FunctionStub* stub) {
 230     DCHECK(stub != NULL);
 231     FunctionStub* to_free = reinterpret_cast<FunctionStub*>(stub);
 232     to_free->FunctionStub::~FunctionStub();
 233     BOOL success = VirtualFree(to_free, sizeof(FunctionStub), MEM_DECOMMIT);
 234     DCHECK(success) << "VirtualFree";
 235     return success != FALSE;
 236   }
 237 };
 238
 239 #pragma pack(pop)
 240
 241 #endif  // CHROME_FRAME_FUNCTION_STUB_H_