SMALL HOOK LIBRARY

#include "winlite.h"
#include "hde32.h"

#ifndef DWORD
typedef unsigned int DWORD;
#endif
#ifndef BYTE
typedef unsigned char BYTE;
#endif

void RemoveHook(DWORD dwOrigFn, DWORD dwHookType)
void InstallHook(DWORD dwOrigFn, DWORD dwHookFn, DWORD dwHookType);

#define XALLOC(sz) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sz)
#define XFREE(ptr) HeapFree(GetProcessHeap(), 0, (ptr))
#define XEXEC(ptr, sz) { DWORD temp; VirtualProtect(ptr, sz, PAGE_EXECUTE_READWRITE, &temp; }

typedef enum {
    HOOK_TYPE_JMP,
    HOOK_TYPE_RET,
} HOOK_TYPES;

void InstallHook(DWORD dwOrigFn, DWORD dwHookFn, DWORD dwHookType)
{
    if (dwHookType == HOOK_TYPE_JMP)
    {
        //use hde32 to calculate trampoline size
        //as to not cut instructions in half
        DWORD dwTrampSz = 0;
        hde32s hde; memset((void*)&hde, 0, sizeof(hde32s));
        while (dwTrampSz < 10)
        {
            hde32_disasm((void*)(dwOrigFn + dwTrampSz), &hde);
            dwTrampSz += hde.len;
        }

        //allocate trampoline function
        //and make executable
        DWORD dwTrampFn = XALLOC(dwTrampSz + 5);
        XEXEC((void*)dwTrampFn, dwTrampSz + 5);

        //the hook code
        BYTE dwHook[10];
        memset(dwHook, 0, 10);
        dwHook[0] = 0xE9;
        *(DWORD*)&dwHook[1] = dwHookFn;
        dwHook[5] = 0x68;
        *(DWORD*)&dwHook[6] = dwTrampFn;

        //copy bytes from the original function
        //to the trampoline function
        memcpy((void*)dwOrigFn, (void*)dwTrampFn, dwTrampSz);

        //and the jump back to the original function
        BYTE dwTrampReturn[5];
        memset(dwTrampReturn, 0, 5);
        dwTrampReturn[0] = 0xE9;
        *(DWORD*)&dwTrampReturn[1] = dwHookFn + dwTrampSz;
        memcpy((void*)&dwTrampReturn[0], (void*)(dwTrampFn + dwTrampSz), 5);

        //install the hook
        memcpy((void*)&dwHook[0], (void*)dwOrigFn, 10);

        //pad with nops if trampoline size
        //is larger than needed
        while (dwTrampSz > 10)
        {
            *(BYTE*)(dwOrigFn + dwTrampSz) = 0x90;
            dwTrampSz--;
        }
    }
    else if (dwHookType == HOOK_TYPE_RET)
    {
        //use hde32 to calculate trampoline size
        //as to not cut instructions in half
        DWORD dwTrampSz = 0;
        hde32s hde; memset((void*)&hde, 0, sizeof(hde32s));
        while (dwTrampSz < 11)
        {
            hde32_disasm((void*)(dwOrigFn + dwTrampSz), &hde);
            dwTrampSz += hde.len;
        }

        //allocate trampoline function
        //and make executable
        DWORD dwTrampFn = XALLOC(dwTrampSz);
        XEXEC((void*)dwTrampFn, dwTrampSz);

        //the hook code
        BYTE dwHook[11];
        memset(dwHook, 0, 11);
        dwHook[0] = 0x68;
        *(DWORD*)&dwHook[1] = dwTrampFn;
        dwHook[5] = 0x68;
        *(DWORD*)&dwHook[6] = dwHookFn;
        *(BYTE*)&dwHook[10] = 0xC3;

        //copy bytes from the original function
        //to the trampoline function
        memcpy((void*)dwOrigFn, (void*)dwTrampFn, dwTrampSz);

        //and the jump back to the original function
        BYTE dwTrampReturn[5];
        memset(dwTrampReturn, 0, 5);
        dwTrampReturn[0] = 0xE9;
        *(DWORD*)&dwTrampReturn[1] = dwHookFn + dwTrampSz;
        memcpy((void*)&dwTrampReturn[0], (void*)(dwTrampFn + dwTrampSz), 5);

        //install the hook
        memcpy((void*)&dwHook[0], (void*)dwOrigFn, 11);

        //pad with nops if trampoline size
        //is larger than needed
        while (dwTrampSz > 11)
        {
            *(BYTE*)(dwOrigFn + dwTrampSz) = 0x90;
            dwTrampSz--;
        }
    }
}

void RemoveHook(DWORD dwOrigFn, DWORD dwHookType)
{
    if (dwHookType == HOOK_TYPE_JMP)
    {
        //get the address of the trampoline function
        DWORD dwTrampFn = *(DWORD*)(dwOrigFn + 1);

        //use hde32 to calculate trampoline size
        DWORD dwTrampSz = 0;
        hde32s hde; memset((void*)&hde, 0, sizeof(hde32s));
        while (dwTrampSz < 10)
        {
            hde32_disasm((void*)(dwTrampFn + dwTrampSz), &hde);
            if (*(DWORD*)(dwTrampFn + dwTrampSz + 1) == dwOrigFn) break;
            dwTrampSz += hde.len;
        }

        //copy bytes from the trampoline function
        //back to the original function (removes the hook)
        memcpy((void*)dwTrampFn, (void*)dwOrigFn, dwTrampSz);

        //free trampoline function
        XFREE(dwTrampSz);
    }
    else if (dwHookType == HOOK_TYPE_RET)
    {
        //get the address of the trampoline function
        DWORD dwTrampFn = *(DWORD*)(dwOrigFn + 6);

        //use hde32 to calculate trampoline size
        DWORD dwTrampSz = 0;
        hde32s hde; memset((void*)&hde, 0, sizeof(hde32s));
        while (dwTrampSz < 11)
        {
            hde32_disasm((void*)(dwTrampFn + dwTrampSz), &hde);
            if (*(DWORD*)(dwTrampFn + dwTrampSz + 1) == dwOrigFn) break;
            dwTrampSz += hde.len;
        }

        //copy bytes from the trampoline function
        //back to the original function (removes the hook)
        memcpy((void*)dwTrampFn, (void*)dwOrigFn, dwTrampSz);

        //free trampoline function
        XFREE(dwTrampSz);
    }
}