VMTHook class x86 und AMD64 kompatibel

#ifndef __VMTHOOK_H__
#define __VMTHOOK_H__
 
#include <Windows.h>
 
class CVMTHook
{
public:
	CVMTHook(void* instance);
	~CVMTHook();
	void* hookFunction(size_t iIndex, void* pfnHook);
	void* getOriginalFunction(size_t iIndex);
	void setHookEnabled(bool bEnabled=true);
protected:
	size_t m_iNumIndices;
	void** m_pOriginalVTable;
	void** m_pNewVTable;
	void*** m_pInstance;
};
 
#endif //__VMTHOOK_H__

#include "VMTHook.h"
 
CVMTHook::CVMTHook(void* instance)
{
	HANDLE hProcessHeap;
 
	if(instance)
	{
		m_pInstance = (void***) instance;
		m_pOriginalVTable = *m_pInstance;
 
		//Count number of Pointers in the table
 
		m_iNumIndices = 0;
 
		//TODO: check if pointer into .text section
		while(m_pOriginalVTable[m_iNumIndices])
		{
			m_iNumIndices++;
		}
 
 
		//Allocate memory on the heap for our own copy of the table
 
		hProcessHeap = GetProcessHeap();
 
		if(hProcessHeap)
		{
			m_pNewVTable = (void**) HeapAlloc(hProcessHeap, 0, sizeof(void*) * m_iNumIndices);
			if(m_pNewVTable)
			{
				memcpy(m_pNewVTable, m_pOriginalVTable, sizeof(void*) * m_iNumIndices);
				setHookEnabled();
			}
		}
	}
}
 
CVMTHook::~CVMTHook()
{
	HANDLE hProcessHeap;
 
	//Reset the VTable pointer
	if(*m_pInstance == m_pNewVTable)
	{
		*m_pInstance = m_pOriginalVTable;
	}
 
	//Free our copy of the VTable
	hProcessHeap = GetProcessHeap();
	if(hProcessHeap)
	{
 
		HeapFree(hProcessHeap, 0, m_pNewVTable);
	}
}
 
 
void* CVMTHook::getOriginalFunction(size_t iIndex)
{
	return m_pOriginalVTable[iIndex];
}
 
void* CVMTHook::hookFunction(size_t iIndex, void* pfnHook)
{
	//Valid index?
	if(iIndex >= m_iNumIndices)
		return NULL;
 
	//Write new pointer
	m_pNewVTable[iIndex]=pfnHook;
 
	//And return pointer to original function
	return m_pOriginalVTable[iIndex];
}
 
void CVMTHook::setHookEnabled(bool bEnabled)
{
	if(bEnabled)
	{
		//Point to our copy of the VTable
		*m_pInstance=m_pNewVTable;
	}
	else
	{
		//Point to the original VTable
		*m_pInstance=m_pOriginalVTable;
	}
}