MinHook - The Minimalistic x86/x64 API Hooking Library

/* 
 *  MinHook - Minimalistic API Hook Library	
 *  Copyright (C) 2009 Tsuda Kageyu. All rights reserved.
 *  
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *  
 *  1. Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *  2. Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *  3. The name of the author may not be used to endorse or promote products
 *     derived from this software without specific prior written permission.
 *  
 *  THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 *  IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 *  OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 *  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 *  NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 *  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 *  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 *  THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <cassert>
#include <vector>
#include <algorithm>
#include <functional>
#include <boost/scope_exit.hpp>
#include <boost/foreach.hpp>
#include <Windows.h>
#include "pstdint.h"

#include "../MinHook.h"
#include "hook.h"
#include "buffer.h"
#include "trampoline.h"
#include "thread.h"

namespace MinHook { namespace
{
	struct HOOK_ENTRY
	{
		void*	pTarget;
		void*	pDetour;
#if defined _M_X64
		void*	pRelay;
#endif
		void*	pTrampoline;
		void*	pBackup;
		bool	isInstalled;
		std::vector<uintptr_t>	oldIPs;
		std::vector<uintptr_t>	newIPs;
	};

	// ���ߏ������ݗp�\����
#pragma pack(push, 1)
	struct JMP_REL
	{
		uint8_t		opcode;
		uint32_t	operand;
	};

	struct JMP_ABS
	{
		uint16_t	opcode;
		uint32_t	operand;
	};
#pragma pack(pop)

	bool		IsHookInstalled(const HOOK_ENTRY& hook);
	HOOK_ENTRY* FindHook(void* const pTarget);
	bool		IsExecutableAddress(void* pAddress);
	void		WriteRelativeJump(void* pFrom, void* const pTo);
	void		WriteAbsoluteJump(void* pFrom, void* const pTo, void* pTable);

	template <typename T>
	bool operator <(const HOOK_ENTRY& lhs, const T& rhs) ;
	template <typename T>
	bool operator <(const T& lhs, const HOOK_ENTRY& rhs) ;
	bool operator <(const HOOK_ENTRY& lhs, const HOOK_ENTRY& rhs);

	CriticalSection gCS;
	std::vector<HOOK_ENTRY> gHooks;
	bool gIsInitialized = false;
}}

namespace MinHook
{
	MH_STATUS Initialize()
	{
		CriticalSection::ScopedLock lock(gCS);

		if (gIsInitialized)
		{
			return MH_ERROR_ALREADY_INITIALIZED;
		}

		// ����֐��o�b�t�@�̏�����
		InitializeBuffer();

		gIsInitialized = true;
		return MH_OK;
	}

	MH_STATUS Uninitialize()
	{
		CriticalSection::ScopedLock lock(gCS);

		if (!gIsInitialized)
		{
			return MH_ERROR_NOT_INITIALIZED;
		}

		// ���ׂẴt�b�N����
		BOOST_FOREACH (const HOOK_ENTRY& hook, gHooks)
		{
			if (!hook.isInstalled)
			{
				continue;
			}
			
			MH_STATUS status = UninstallHook(hook.pTarget);
			if (status != MH_OK)
			{
				return status;
			}
		}

		std::vector<HOOK_ENTRY> v;
		gHooks.swap(v);

		// ����֐��o�b�t�@�̊J��
		UninitializeBuffer();

		gIsInitialized = false;
		return MH_OK;
	}

	MH_STATUS InstallHook(void* pTarget, void* const pDetour)
	{
		CriticalSection::ScopedLock lock(gCS);
		if (!IsExecutableAddress(pTarget) || !IsExecutableAddress(pDetour))
		{
			return MH_ERROR_NOT_EXECUTABLE;
		}

		HOOK_ENTRY *pHook = FindHook(pTarget);
		if (pHook != NULL && pHook->isInstalled)
		{
			return MH_ERROR_ALREADY_INSTALLED;
		}

		if (pHook == NULL)
		{
			bool committed = false;
			BOOST_SCOPE_EXIT((&committed))
			{
				if (!committed)
				{
					RollbackBuffer();
				}
			}
			BOOST_SCOPE_EXIT_END;

			// �g�����|�����֐���쐬����
			CREATE_TREMPOLINE_T ct = { 0 };
			ct.pTarget = pTarget;
			if (!CreateTrampolineFunction(ct))
			{
				return MH_ERROR_UNSUPPORTED_FUNCTION;
			}

			void* pTrampoline = AllocateCodeBuffer(pTarget, ct.trampoline.size());
			if (pTrampoline == NULL)
			{
				return MH_ERROR_MEMORY_ALLOC;
			}
#if defined _M_X64
			void* pTable = AllocateDataBuffer(pTrampoline, (ct.table.size() + 1) * sizeof(uintptr_t));
			if (pTable == NULL)
			{
				return MH_ERROR_MEMORY_ALLOC;
			}
#endif

			ct.pTrampoline = pTrampoline;
#if defined _M_X64
			ct.pTable = pTable;
#endif
			if (!ResolveTemporaryAddresses(ct))
			{
				return MH_ERROR_UNSUPPORTED_FUNCTION;
			}

			memcpy(pTrampoline, &ct.trampoline[ 0 ], ct.trampoline.size());
#if defined _M_X64
			if (ct.table.size() != 0)
			{
				memcpy(pTable, &ct.table[ 0 ], ct.table.size() * sizeof(uintptr_t));
			}
#endif

			// �^�[�Q�b�g�֐��̃o�b�N�A�b�v��Ƃ�
			void* pBackup = AllocateDataBuffer(NULL, sizeof(JMP_REL));
			if (pBackup == NULL)
			{
				return MH_ERROR_MEMORY_ALLOC;
			}

			memcpy(pBackup, pTarget, sizeof(JMP_REL));

			// ���p�֐���쐬����
#if defined _M_X64
			void* pRelay = AllocateCodeBuffer(pTarget, sizeof(JMP_ABS));
			if (pRelay == NULL)
			{
				return MH_ERROR_MEMORY_ALLOC;
			}

			WriteAbsoluteJump(pRelay, pDetour, reinterpret_cast<uintptr_t*>(pTable) + ct.table.size());
#endif
			CommitBuffer();
			committed = true;

			// �t�b�N���̓o�^
			HOOK_ENTRY hook = { 0 };
			hook.pTarget = pTarget;
			hook.pDetour = pDetour;
#if defined _M_X64
			hook.pRelay  = pRelay;
#endif
			hook.pTrampoline = pTrampoline;
			hook.pBackup = pBackup;
			hook.oldIPs = ct.oldIPs;
			hook.newIPs = ct.newIPs;

			std::vector<HOOK_ENTRY>::iterator i	= std::lower_bound(gHooks.begin(), gHooks.end(), hook);
			i = gHooks.insert(i, hook);
			pHook = &(*i);

		}

		// �^�[�Q�b�g�֐��̖`���ɁA���p�֐��܂��̓t�b�N�֐��ւ̃W�����v���������
		{
			ScopedThreadExclusive tex(pHook->oldIPs, pHook->newIPs);

			DWORD oldProtect;
			if (!VirtualProtect(pHook->pTarget, sizeof(JMP_REL), PAGE_EXECUTE_READWRITE, &oldProtect))
			{
				return MH_ERROR_MEMORY_PROTECT;
			}

#if defined _M_X64
			WriteRelativeJump(pHook->pTarget, pHook->pRelay);
#elif defined _M_IX86
			WriteRelativeJump(pHook->pTarget, pHook->pDetour);
#endif
			VirtualProtect(pHook->pTarget, sizeof(JMP_REL), oldProtect, &oldProtect);
		}

		pHook->isInstalled = true;
		return MH_OK;
	}

	MH_STATUS UninstallHook(void* pTarget)
	{
		CriticalSection::ScopedLock lock(gCS);

		if (!IsExecutableAddress(pTarget))
		{
			return MH_ERROR_NOT_EXECUTABLE;
		}

		HOOK_ENTRY *pHook = FindHook(pTarget);
		if (pHook == NULL || !pHook->isInstalled)
		{
			return MH_ERROR_NOT_INSTALLED;
		}

		// �^�[�Q�b�g�֐��̖`��������߂������B���͍ė��p�̂��ߎc���Ă���
		{
			ScopedThreadExclusive tex(pHook->oldIPs, pHook->newIPs);

			DWORD oldProtect;
			if (!VirtualProtect(pHook->pTarget, sizeof(JMP_REL), PAGE_EXECUTE_READWRITE, &oldProtect))
			{
				return MH_ERROR_MEMORY_PROTECT;
			}

			memcpy(pHook->pTarget, pHook->pBackup, sizeof(JMP_REL));

			VirtualProtect(pHook->pTarget, sizeof(JMP_REL), oldProtect, &oldProtect);
		}

		pHook->isInstalled = false;
		return MH_OK;
	}

	MH_STATUS GetOriginalFunction(void* pTarget, void*& pTrampoline)
	{
		CriticalSection::ScopedLock lock(gCS);

		HOOK_ENTRY *pHook = FindHook(pTarget);
		if (pHook == NULL || !pHook->isInstalled)
		{
			return MH_ERROR_NOT_INSTALLED;
		}

		pTrampoline = pHook->pTrampoline;
		return MH_OK;
	}

	MH_STATUS GetHookCount(size_t& count)
	{
		using std::tr1::bind;
		using std::tr1::placeholders::_1;

		CriticalSection::ScopedLock lock(gCS);

		count = std::count_if(gHooks.begin(), gHooks.end(), bind(IsHookInstalled, _1));
		return MH_OK;
	}
}
namespace MinHook { namespace
{
	bool IsHookInstalled(const HOOK_ENTRY& hook)
	{
		return (hook.isInstalled);
	}

	HOOK_ENTRY* FindHook(void* const pTarget)
	{
		std::vector<HOOK_ENTRY>::iterator i 
			= std::lower_bound(gHooks.begin(), gHooks.end(), pTarget);
		if (i != gHooks.end() && i->pTarget == pTarget)
		{
			return &(*i);
		}

		return NULL;
	}

	bool IsExecutableAddress(void* pAddress)
	{
		static const DWORD PageExecuteMask 
			= (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY); 

		// �����蓖�Ă���s�s�”\�ȗ̈��`�F�b�N
		MEMORY_BASIC_INFORMATION mi = { 0 };
		VirtualQuery(pAddress, &mi, sizeof(mi));

		return ((mi.Protect & PageExecuteMask) != 0);
	}

	void WriteRelativeJump(void* pFrom, void* const pTo)
	{
		JMP_REL *pInst = reinterpret_cast<JMP_REL*>(pFrom);
		pInst->opcode = 0xE9;
		pInst->operand = static_cast<uint32_t>(
			reinterpret_cast<char*>(pTo) - (reinterpret_cast<char*>(pFrom) + sizeof(JMP_REL)));
	}

	void WriteAbsoluteJump(void* pFrom, void* const pTo, void* pTable)
	{
		JMP_ABS *pInst = reinterpret_cast<JMP_ABS*>(pFrom);
		pInst->opcode = 0x25FF;
		pInst->operand = static_cast<uint32_t>(
			reinterpret_cast<char*>(pTable) - (reinterpret_cast<char*>(pFrom) + sizeof(JMP_ABS)));

		void** pAddr = reinterpret_cast<void**>(pTable);
		*pAddr = pTo;
	}



	template <typename T>
	bool operator <(const HOOK_ENTRY& lhs, const T& rhs) 
	{ 
		return lhs.pTarget < reinterpret_cast<void*>(rhs); 
	}

	template <typename T>
	bool operator <(const T& lhs, const HOOK_ENTRY& rhs) 
	{ 
		return reinterpret_cast<void*>(lhs) < rhs.pTarget; 
	}

	bool operator <(const HOOK_ENTRY& lhs, const HOOK_ENTRY& rhs)
	{ 
		return lhs.pTarget < rhs.pTarget;
	}
}}