#include <windows.h>
#include "RealAddress.h"


FARPROC RealAddress::GetProcAddress(HMODULE hModule, LPCSTR lpProcName)
{
	FARPROC pResolved2, pResolved1 = ::GetProcAddress(hModule, lpProcName);

	if(0 == pResolved1)
	{
		return NULL;
	}

	if(FALSE == this->m_bUsingThunks)
	{
		pResolved2 = ::GetProcAddress(hModule, lpProcName);

		if(0 == pResolved2)
		{
			// DLL was unloaded from another thread
			return NULL;
		}

		if(pResolved1 == pResolved2)
		{
			// all is good. return the address
			return pResolved1;
		}
	}

	PBYTE pDebugThunk = reinterpret_cast<PBYTE>(pResolved1);

	//
	// 0x68, ??, ??, ??, ??, 0xE9, ??, ??, ??, ??
	//

	if(0x68 != *pDebugThunk || 0xE9 != *(pDebugThunk + 5))
	{
		// unknown debug thunk format
		return pResolved1;
	}

	this->m_nJumpTarget = reinterpret_cast<DWORD>(pDebugThunk + 5 + 5) +
		*(reinterpret_cast<PDWORD>(pDebugThunk + 6));

	this->m_bUsingThunks = TRUE;

	// get the real address from debug thunk
	return *(reinterpret_cast<FARPROC *>(pDebugThunk + 1));
}


DWORD RealAddress::GetThunkAddress(PIMAGE_THUNK_DATA32 pThunk)
{
	if(NULL == pThunk || NULL == pThunk->u1.Function)
	{
		return NULL;
	}

/*
	if(0 == (pThunk->u1.Function & 0x80000000))
	{
		// thunks are put above 2 GB
		return pThunk->u1.Function;
	}
*/

	if(FALSE == this->m_bUsingThunks)
	{
		HMODULE hKernel32 = ::GetModuleHandle(TEXT("kernel32.dll"));

		FARPROC pResolved1 = ::GetProcAddress(hKernel32, "GetProcAddress");
		FARPROC pResolved2 = ::GetProcAddress(hKernel32, "GetProcAddress");

		if(pResolved1 == pResolved2)
		{
			return pThunk->u1.Function;
		}

		PBYTE pDebugThunk = reinterpret_cast<PBYTE>(pResolved1);

		if(0x68 != *pDebugThunk || 0xE9 != *(pDebugThunk + 5))
		{
			// rats. unknown format
			return pThunk->u1.Function;
		}

		this->m_nJumpTarget = reinterpret_cast<DWORD>(pDebugThunk + 5 + 5) +
			*(reinterpret_cast<PDWORD>(pDebugThunk + 6));

		this->m_bUsingThunks = TRUE;
	}

	//
	// start working on the import table
	//

	PBYTE pDebugThunk = reinterpret_cast<PBYTE>(pThunk->u1.Function);

	if(0x68 != *pDebugThunk || 0xE9 != *(pDebugThunk + 5))
	{
		// probably not a debug thunk
		return pThunk->u1.Function;
	}

	DWORD nJumpTarget = reinterpret_cast<DWORD>(pDebugThunk + 5 + 5) +
		*(reinterpret_cast<PDWORD>(pDebugThunk + 6));

	if(nJumpTarget != this->m_nJumpTarget)
	{
		// evidently not a debug thunk
		return pThunk->u1.Function;
	}

	// get the real address from thunk
	return *(reinterpret_cast<PDWORD>(pDebugThunk + 1));
}