// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.


#ifndef CHROME_FRAME_FUNCTION_STUB_H_
#define CHROME_FRAME_FUNCTION_STUB_H_

#include <windows.h>
#include "base/logging.h"

// IMPORTANT: The struct below must be byte aligned.
#pragma pack(push)
#pragma pack(1)

#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
// for patching entries of a vtable or e.g. a globally registered wndproc
// for a class as opposed to a window.
// When unhooking, you can just call the BypassStub() function and leave the
// stub in memory.  This unhooks your function while leaving the (potential)
// chain of patches intact.
//
// @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 
//        operation and then restore it.
//
// For instance if the function prototype is:
//
// @code
//   LRESULT WndProc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam);
// @endcode
//
// and we would like to add one static argument to make it, say:
//
// @code
//   LRESULT MyNewWndProc(WNDPROC original, HWND hwnd, UINT msg, 
//                        WPARAM wparam, LPARAM lparam);
// @endcode
//
// That can be achieved by wrapping the function up with a FunctionStub:
//
// @code
//   FunctionStub* stub = FunctionStub::Create(original_wndproc, MyNewWndProc);
//   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);

    // 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;
  }

  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;
  }

  inline PROC code() const {
    return reinterpret_cast<PROC>(const_cast<Opcodes*>(&code_));
  }

  // 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);

    // 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));

    if (!ret) {
      NOTREACHED();
    } else {
      // Construct
      ret->FunctionStub::FunctionStub(extra_argument, dest);
    }

    return ret;
  }

  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;
  }

  // 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;
  }
};

#pragma pack(pop)

#endif  // CHROME_FRAME_FUNCTION_STUB_H_