Reimplementation of FunctionStub, to avoid rewriting potentially executing co...

chrome_frame/function_stub.h

Index: chrome_frame/function_stub.h
===================================================================
--- chrome_frame/function_stub.h	(revision 42235)
+++ chrome_frame/function_stub.h	(working copy)
@@ -13,12 +13,28 @@
 #pragma pack(push)
 #pragma pack(1)
 
+struct FunctionStubAsm {
+  // The stub always starts with an indirect jump, which starts out jumping
+  // to the remainder of the stub. This means we can bypass the stub by
+  // rewriting the jump destination, which is data, in a manner that does
+  // not involve writing code, only writing data at a natural word boundary.
+  uint16 jump_to_bypass_;        // indirect jump
+  uintptr_t bypass_target_addr_; // to the bypass target.
+  uint8 pop_return_addr_;        // pop eax
+  uint16 push_;                  // push [arg]   ; push...
+  uintptr_t arg_addr_;           //              ; extra argument
+  uint8 push_return_addr_;       // push eax     ; push the return address
+  uint16 jump_to_target;         // jmp [target] ; jump...
+  uintptr_t target_addr_;        //              ; to the hook function
+};
+
+#pragma pack(pop)
+
+
 #ifndef _M_IX86
 #error Only x86 supported right now.
 #endif
 
-extern "C" IMAGE_DOS_HEADER __ImageBase;
-
 // This struct is assembly code + signature.  The purpose of the struct is to be
 // able to hook an existing function with our own and store information such
 // as the original function pointer with the code stub.  Typically this is used
@@ -30,7 +46,7 @@
 //
 // @note: This class is meant for __stdcall calling convention and
 //        it uses eax as a temporary variable.  The struct can
-//        be improved in the future to save eax before the 
+//        be improved in the future to save eax before the
 //        operation and then restore it.
 //
 // For instance if the function prototype is:
@@ -42,7 +58,7 @@
 // and we would like to add one static argument to make it, say:
 //
 // @code
-//   LRESULT MyNewWndProc(WNDPROC original, HWND hwnd, UINT msg, 
+//   LRESULT MyNewWndProc(WNDPROC original, HWND hwnd, UINT msg,
 //                        WPARAM wparam, LPARAM lparam);
 // @endcode
 //
@@ -53,189 +69,73 @@
 //   SetClassLongPtr(wnd, GCLP_WNDPROC, stub->code());
 // @endcode
 struct FunctionStub {
- private:
-  typedef enum AsmConstants {
-    POP_EAX = 0x58,
-    PUSH = 0x68,
-    PUSH_EAX = 0x50,
-    JUMP_RELATIVE = 0xE9
-  };
-
-  FunctionStub(uintptr_t extra_argument, void* dest)
-      : signature_(reinterpret_cast<HMODULE>(&__ImageBase)) {
-    Opcodes::Hook& hook = code_.hook_;
-    hook.pop_return_addr_ = POP_EAX;
-    hook.push_ = PUSH;
-    hook.arg_ = extra_argument;
-    hook.push_return_addr_ = PUSH_EAX;
-    hook.jump_to_ = JUMP_RELATIVE;
-
-    // Calculate the target jump to the destination function.
-    hook.target_ = CalculateRelativeJump(dest, &hook.jump_to_);
-
-    // Allow the process to execute this struct as code.
-    DWORD old_protect = 0;
-    // Allow reads too since we want read-only member variable access at
-    // all times.
-    ::VirtualProtect(this, sizeof(FunctionStub), PAGE_EXECUTE_READ,
-                     &old_protect);
-  }
-
-  ~FunctionStub() {
-    // No more execution rights.
-    DWORD old_protect = 0;
-    ::VirtualProtect(this, sizeof(FunctionStub), PAGE_READWRITE, &old_protect);
-  }
-
-  // Calculates the target value for a relative jump.
-  // The function assumes that the size of the opcode is 1 byte.
-  inline uintptr_t CalculateRelativeJump(const void* absolute_to,
-                                         const void* absolute_from) const {
-    return (reinterpret_cast<uintptr_t>(absolute_to) -
-            reinterpret_cast<uintptr_t>(absolute_from)) - 
-           (sizeof(uint8) + sizeof(uintptr_t));
-  }
-
-  // Does the opposite of what CalculateRelativeJump does, which is to
-  // calculate back the absolute address that previously was relative to
-  // some other address.
-  inline uintptr_t CalculateAbsoluteAddress(const void* relative_to,
-                                            uintptr_t relative_address) const {
-    return relative_address + sizeof(uint8) + sizeof(uintptr_t) +
-           reinterpret_cast<uintptr_t>(relative_to);
-  }
-
-  // Used to identify function stubs that belong to this module.
-  HMODULE signature_;
-
-  // IMPORTANT: Do not change the order of member variables
-  union Opcodes {
-    // Use this struct when the stub forwards the call to our hook function
-    // providing an extra argument.
-    struct Hook {
-      uint8 pop_return_addr_;  // pop eax
-      uint8 push_;             // push arg  ; push...
-      uintptr_t arg_;          //           ; extra argument
-      uint8 push_return_addr_; // push eax  ; push the return address
-      uint8 jump_to_;          // jmp       ; jump...
-      uintptr_t target_;       //           ; to the hook function
-    } hook_;
-    // When the stub is bypassed, we jump directly to a given target,
-    // usually the originally hooked function.
-    struct Bypassed {
-      uint8 jump_to_;          // jmp to
-      uintptr_t target_;       // relative target.
-    } bypassed_;
-  };
-
-  Opcodes code_;
-
  public:
   // Neutralizes this stub and converts it to a direct jump to a new target.
-  void BypassStub(void* new_target) {
-    DWORD old_protect = 0;
-    // Temporarily allow us to write to member variables
-    ::VirtualProtect(this, sizeof(FunctionStub), PAGE_EXECUTE_READWRITE,
-                     &old_protect);
+  void BypassStub(void* new_target);
 
-    // Now, just change the first 5 bytes to jump directly to the new target.
-    Opcodes::Bypassed& bypassed = code_.bypassed_;
-    bypassed.jump_to_ = JUMP_RELATIVE;
-    bypassed.target_ = CalculateRelativeJump(new_target, &bypassed.jump_to_);
-
-    // Restore the previous protection flags.
-    ::VirtualProtect(this, sizeof(FunctionStub), old_protect, &old_protect);
-
-    // Flush the instruction cache just in case.
-    ::FlushInstructionCache(::GetCurrentProcess(), this, sizeof(FunctionStub));
-  }
-
-  // @returns the argument supplied in the call to @ref Create
-  inline uintptr_t argument() const {
-    DCHECK(code_.hook_.pop_return_addr_ == POP_EAX) << "stub has been disabled";
-    return code_.hook_.arg_;
-  }
-
   inline bool is_bypassed() const {
-    return code_.bypassed_.jump_to_ == JUMP_RELATIVE;
+    return bypass_address_ !=
+        reinterpret_cast<uintptr_t>(&stub_.pop_return_addr_);
   }
 
-  inline uintptr_t absolute_target() const {
-    DCHECK(code_.hook_.pop_return_addr_ == POP_EAX) << "stub has been disabled";
-    return CalculateAbsoluteAddress(
-        reinterpret_cast<const void*>(&code_.hook_.jump_to_),
-        code_.hook_.target_);
-  }
-
   // Returns true if the stub is valid and enabled.
   // Don't call this method after bypassing the stub.
-  inline bool is_valid() const {
-    return signature_ == reinterpret_cast<HMODULE>(&__ImageBase) &&
-        code_.hook_.pop_return_addr_ == POP_EAX;
-  }
+  bool is_valid() const;
 
   inline PROC code() const {
-    return reinterpret_cast<PROC>(const_cast<Opcodes*>(&code_));
+    return reinterpret_cast<PROC>(const_cast<FunctionStubAsm*>(&stub_));
   }
 
   // Use to create a new function stub as shown above.
-  //
   // @param extra_argument The static argument to pass to the function.
   // @param dest Target function to which the stub applies.
   // @returns NULL if an error occurs, otherwise a pointer to the
-  // function stub.
-  //
-  // TODO(tommi): Change this so that VirtualAlloc isn't called for
-  // every stub.  Instead we should re-use each allocation for
-  // multiple stubs.  In practice I'm guessing that there would
-  // only be one allocation per process, since each allocation actually
-  // allocates at least one page of memory (4K).  Size of FunctionStub
-  // is 12 bytes, so one page could house 341 function stubs.
-  // When stubs are created frequently, VirtualAlloc could fail
-  // and last error is ERROR_NOT_ENOUGH_MEMORY (8).
-  static FunctionStub* Create(uintptr_t extra_argument, void* dest) {
-    DCHECK(dest);
+  //    function stub.
+  static FunctionStub* Create(uintptr_t extra_argument, void* dest);
 
-    // Use VirtualAlloc to get memory for the stub.  This gives us a
-    // whole page that we don't share with anyone else.
-    // Initially the memory must be READWRITE to allow for construction
-    // PAGE_EXECUTE is set in constructor.
-    FunctionStub* ret = reinterpret_cast<FunctionStub*>(VirtualAlloc(NULL, 
-        sizeof(FunctionStub), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
+  // Test whether address (likely) points to an existing function stub.
+  // @returns NULL if address does not point to a function stub.
+  // @note likely means approximately 1/2^48 here.
+  static FunctionStub* FromCode(void* address);
 
-    if (!ret) {
-      NOTREACHED();
-    } else {
-      // Construct
-      ret->FunctionStub::FunctionStub(extra_argument, dest);
-    }
+  // Deallocates a FunctionStub.
+  // The stub must not be in use on any thread!
+  static bool Destroy(FunctionStub* stub);
 
-    return ret;
-  }
+  // Accessors.
+  uintptr_t argument() const { return argument_; }
+  void set_argument(uintptr_t argument) { argument_ = argument; }
 
-  static FunctionStub* FromCode(void* address) {
-    Opcodes* code = reinterpret_cast<Opcodes*>(address);
-    if (code->hook_.pop_return_addr_ == POP_EAX) {
-      FunctionStub* stub = reinterpret_cast<FunctionStub*>(
-          reinterpret_cast<uint8*>(address) - sizeof(HMODULE));
-      if (stub->is_valid())
-        return stub;
-    }
-
-    return NULL;
+  uintptr_t bypass_address() const { return bypass_address_; }
+  void set_bypass_address(uintptr_t bypass_address) {
+    bypass_address_ = bypass_address;
   }
 
-  // Deallocates a FunctionStub.  The stub must not be in use on any thread!
-  static bool Destroy(FunctionStub* stub) {
-    DCHECK(stub != NULL);
-    FunctionStub* to_free = reinterpret_cast<FunctionStub*>(stub);
-    to_free->FunctionStub::~FunctionStub();
-    BOOL success = VirtualFree(to_free, sizeof(FunctionStub), MEM_DECOMMIT);
-    DCHECK(success) << "VirtualFree";
-    return success != FALSE;
+  uintptr_t destination_function() const { return destination_function_; }
+  void set_destination_function(uintptr_t destination_function) {
+    destination_function_ = destination_function;
   }
-};
 
-#pragma pack(pop)
+ protected:
+  // Protected for testing only.
+  FunctionStub(uintptr_t extra_argument, void* dest);
+  ~FunctionStub();
 
+  void Init(FunctionStubAsm* stub);
+
+  FunctionStubAsm stub_;
+
+  // Used to identify function stubs that belong to this module.
+  HMODULE signature_;
+
+  // This is the argument value that gets passed to the destination_function_.
+  uintptr_t argument_;
+  // Bypass address, if this is the address of the pop_return_addr_, the
+  // function stub is not bypassed.
+  uintptr_t bypass_address_;
+  // The destination function we dispatch to, not used if the stub
+  // is bypassed.
+  uintptr_t destination_function_;
+};
+
 #endif  // CHROME_FRAME_FUNCTION_STUB_H_