My new iteration on a Windows-based function hooking class, this time templated!

func_hook.hxx

#ifndef SHARED_FUNC_HOOK_H_
#define SHARED_FUNC_HOOK_H_

#include <array>
#include <Windows.h>
#include "./types.h"

namespace sbat {
// Type for simpler VirtualProtect -> restore old protection usage.
class ScopedVirtualProtect {
public:
  ScopedVirtualProtect(void* address, size_t size, u4 new_protection)
      : address_(address),
        size_(size),
        old_protection_(0),
        has_errors_(false) {
    has_errors_ = VirtualProtect(address_, size_, new_protection,
        reinterpret_cast<PDWORD>(&old_protection_)) == 0;
  }

  ~ScopedVirtualProtect() {
    VirtualProtect(address_, size_, old_protection_, reinterpret_cast<PDWORD>(&old_protection_));
  }

  bool has_errors() const {
    return has_errors_;
  }
private:
  void* address_;
  size_t size_;
  u4 old_protection_;
  bool has_errors_;
};

// Type for hooking a function at a specific memory location, with methods for replacing and
// restoring the original code. Function pointer type is specified by F, to allow for hooks of
// varying parameter lists.
template<typename F>
class FuncHook {
public:
  FuncHook(F func, F hook_func)
      : callable_(func),
        hook_func_(hook_func),
        function_(reinterpret_cast<byte*>(func)),
        original_mem_(),
        hooked_mem_(),
        injected_(false) {
    LoadFunctionMemory();
  }

  ~FuncHook() {
    if (injected_) {
      Restore();
    }
  }

  bool Inject() {
    if (injected_) return false;

    ScopedVirtualProtect protect(function_, original_mem_.size(), PAGE_EXECUTE_READWRITE);
    if (protect.has_errors()) return false;

    for (size_t i = 0; i < hooked_mem_.size(); i++) {
      function_[i] = hooked_mem_[i];
    }
    injected_ = true;
    return true;
  }

  bool Restore() {
    if (!injected_) return false;

    ScopedVirtualProtect protect(function_, original_mem_.size(), PAGE_EXECUTE_READWRITE);
    if (protect.has_errors()) return false;

    for (size_t i = 0; i < original_mem_.size(); i++) {
      function_[i] = original_mem_[i];
    }
    injected_ = false;
    return true;
  }

  F callable() const {
    return callable_;
  }

private:
  void LoadFunctionMemory() {
    ScopedVirtualProtect protect(function_, original_mem_.size(), PAGE_EXECUTE_READ);

    // get a pointer to the address pointer at the second byte of hooked_mem_ (param for push)
    F* ret_target_ptr = reinterpret_cast<F*>(&hooked_mem_[1]);
    *ret_target_ptr = hook_func_;  // set parameter of push to the address of our hook function
    hooked_mem_[0] = 0x68;  // push (address provided through hookFunc)
    hooked_mem_[5] = 0xc3;  // return

    for (size_t i = 0; i < original_mem_.size(); i++) {
      original_mem_[i] = function_[i];
    }
  }

  F callable_;
  F hook_func_;
  byte* function_;
  std::array<byte,6> original_mem_;
  std::array<byte,6> hooked_mem_;
  bool injected_;
};
}  // namespace sbat
#endif  // SHARED_FUNC_HOOK_H_

old_FuncHook.cxx

#include "FuncHook.h"
#include "ods.h"

FuncHook::FuncHook() {
  funcMemory = NULL;
	injected = false;
	int i;
	for(i = 0; i < 6; i++)
		origMem[i] = newMem[i] = 0;
}

FuncHook::FuncHook(FuncPtr func, FuncPtr hookFunc) {
	funcMemory = NULL;
	injected = false;
	int i;
	for(i = 0; i < 6; i++)
		origMem[i] = newMem[i] = 0;
	setup(func, hookFunc);
}

void FuncHook::setup(FuncPtr func, FuncPtr hookFunc) {
	if(funcMemory != NULL) {// this hook has been setup before
		if(call == func)
			return; // same function, just leave it in place
		else {
			if(injected)
				restore();
		}
	}
	
	injected = false;
	call = func;
	funcMemory = reinterpret_cast<unsigned char *>(func);
	unsigned char *hookFuncMem = reinterpret_cast<unsigned char *>(hookFunc);

	DWORD oldProtect;
	if(VirtualProtect(funcMemory, 6, PAGE_EXECUTE_READ, &oldProtect) == FALSE) return;
	
	unsigned char **newMemAddrPtr = reinterpret_cast<unsigned char **>(&newMem[1]); // get a pointer to the address pointer at the second byte of newMem (param for push)
	*newMemAddrPtr = hookFuncMem; // set parameter of push to the address of our hookfunc
	newMem[0] = 0x68; // push (address provided through hookFunc)
	newMem[5] = 0xc3; // return

	for(int i = 0; i < 6; i++)
		origMem[i] = funcMemory[i];

	//ods("OrigMem at 0x%08x: %02x %02x %02x %02x %02x %02x", call, origMem[0], origMem[1], origMem[2], origMem[3], origMem[4], origMem[5]);

	VirtualProtect(funcMemory, 6, oldProtect, &oldProtect);
}

bool FuncHook::inject() {
	DWORD oldProtect;
	if(VirtualProtect(funcMemory, 6, PAGE_EXECUTE_READWRITE, &oldProtect) == FALSE) return false;
//	ods("Redirecting function at 0x%08x to 0x%08x", call, *(reinterpret_cast<unsigned char **>(&newMem[1])));
	for(int i = 0; i < 6; i++)
		funcMemory[i] = newMem[i];
	injected = true;
	if(VirtualProtect(funcMemory, 6, oldProtect, &oldProtect) == FALSE) return false;
	FlushInstructionCache(GetCurrentProcess(), funcMemory, 6);

	for(int i = 0; i < 6; i++) {
		if(funcMemory[i] != newMem[i])
			ods("inject() failure for 0x%08x at byte %d", call, i);
	}

	return true;
}

bool FuncHook::restore() {
	DWORD oldProtect;
	if(VirtualProtect(funcMemory, 6, PAGE_EXECUTE_READWRITE, &oldProtect) == FALSE) return false;
	for(int i = 0; i < 6; i++)
		funcMemory[i] = origMem[i];
	injected = false;
	if(VirtualProtect(funcMemory, 6, oldProtect, &oldProtect) == FALSE) return false;
	FlushInstructionCache(GetCurrentProcess(), funcMemory, 6);
	
	for(int i = 0; i < 6; i++) {
		if(funcMemory[i] != origMem[i])
			ods("inject() failure for 0x%08x at byte %d", call, i);
	}

	return true;
}

old_FuncHook.hxx

#ifndef FUNCHOOK_H
#define FUNCHOOK_H

#define WIN32_LEAN_AND_MEAN
#include <windows.h>

#define VTABLE_OFFSET(pClass, nFuncIndex) (reinterpret_cast<FuncPtr>(((reinterpret_cast<void ***>(pClass))[0])[nFuncIndex]))


typedef void *(*FuncPtr)();

struct FuncHook {
  unsigned char * funcMemory;
	unsigned char origMem[6];
	unsigned char newMem[6];

	FuncPtr call;

	bool injected;

	FuncHook();
	FuncHook(FuncPtr func, FuncPtr hookFunc);
	void setup(FuncPtr func, FuncPtr hookFunc);
	bool inject();
	bool restore();
};

#endif