root/lang/objective-cplusplus/i3/trunk/tmp/dwmedit/ApiHook.cc @ 39066

#include "ApiHook.h"
#include "WindowsVersionHelp.h"
#include <tlhelp32.h>

namespace api_hook {

    /**
     * �ЂƂ‚̃��W���[���ɑ΂���PI�t�b�N����֐�
     * @see http://ruffnex.oc.to/kenji/text/api_hook/
     */
    void ReplaceIATEntryInOneMod(
        const char* pszModuleName,
        void* pfnCurrent,
        void* pfnNew,
        HMODULE hmodCaller)
    {
        ULONG ulSize = 0;
        IMAGE_IMPORT_DESCRIPTOR* pImportDesc = 
            reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>(
                dll.ImageDirectoryEntryToData(
                    hmodCaller, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &ulSize));

        if (pImportDesc == NULL) {
            return;
        }

        for (; pImportDesc->Name; ++pImportDesc) {
            char* pszModName = reinterpret_cast<char*>(hmodCaller) + pImportDesc->Name;
            if (lstrcmpiA(pszModName, pszModuleName) == 0) {
                break;
            }
        }

        if (!pImportDesc->Name) {
            return;
        }

        IMAGE_THUNK_DATA* pThunk = 
            reinterpret_cast<IMAGE_THUNK_DATA*>(
                reinterpret_cast<char*>(hmodCaller) + pImportDesc->FirstThunk);

        for (; pThunk->u1.Function; ++pThunk) {
            void** ppfn = reinterpret_cast<void**>(&pThunk->u1.Function);
            if (*ppfn != pfnCurrent) {
                continue;
            }
            DWORD dwDummy = 0;
            VirtualProtect(ppfn, sizeof(ppfn), PAGE_EXECUTE_READWRITE, &dwDummy);
            // DLL code will be CoW
            WriteProcessMemory(GetCurrentProcess(), ppfn, &pfnNew, sizeof(pfnNew), NULL); 
            return;
        }
    }

    /**
     * ���ׂẴ��W���[���ɑ΂���PI�t�b�N����֐�
     * @see http://ruffnex.oc.to/kenji/text/api_hook/
     */
    void ReplaceIATEntryInAllMods(
        const char* pszModuleName,
        void* pfnCurrent,
        void* pfnNew)
    {
        // ���W���[�����X�g���
        HANDLE hModuleSnap = CreateToolhelp32Snapshot(
                                 TH32CS_SNAPMODULE, GetCurrentProcessId());
        if(hModuleSnap == INVALID_HANDLE_VALUE) {
            return;
        }
    
        MODULEENTRY32 me = {};
        me.dwSize = sizeof(me);
        BOOL bModuleResult = Module32First(hModuleSnap, &me);
        // ���ꂼ�����W���[���ɑ΂���eplaceIATEntryInOneMod��s
        while (bModuleResult) {
            ReplaceIATEntryInOneMod(pszModuleName, pfnCurrent, pfnNew, me.hModule);
            bModuleResult = Module32Next(hModuleSnap, &me);
        }
        CloseHandle(hModuleSnap);
    }
}