trunk/TTProxy/Hooker.h

#ifndef _YEBISOCKS_HOOKER_H_
#define _YEBISOCKS_HOOKER_H_

#include <YCL/Hashtable.h>
using namespace yebisuya;

class Hooker {
private:
    template<class TYPE>
        class Ptr {
        public:
                static TYPE* cast(void* p, int offset) {
                        return (TYPE*) ((LPBYTE) p + offset);
                }
    };
    static bool modifyModuleProc(HMODULE module, FARPROC oldproc, FARPROC newproc) {
                IMAGE_DOS_HEADER* dos = (IMAGE_DOS_HEADER*) module;
                if (::IsBadReadPtr(dos, sizeof *dos) || dos->e_magic != IMAGE_DOS_SIGNATURE) {
                        return false;
                }
                
                IMAGE_NT_HEADERS* nt = Ptr<IMAGE_NT_HEADERS>::cast(dos, dos->e_lfanew);
                if (::IsBadReadPtr(nt, sizeof *nt) || nt->Signature != IMAGE_NT_SIGNATURE) {
                        return false;
                }
                
                IMAGE_IMPORT_DESCRIPTOR* desc = Ptr<IMAGE_IMPORT_DESCRIPTOR>::cast(module, nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
                if (desc == (IMAGE_IMPORT_DESCRIPTOR*) nt) {
                        return false;
                }
                
                DWORD o;
                for (; desc->Name; desc++){
                        for (IMAGE_THUNK_DATA* thunk = Ptr<IMAGE_THUNK_DATA>::cast(module, desc->FirstThunk); thunk->u1.Function; thunk++){
                                if ((LPDWORD) thunk->u1.Function == (LPDWORD) oldproc) {
                                        if (::VirtualProtect(&thunk->u1.Function, sizeof thunk->u1.Function, PAGE_READWRITE, &o)){
                                                (LPDWORD&) thunk->u1.Function = (LPDWORD) newproc;
                                                ::VirtualProtect(&thunk->u1.Function, sizeof thunk->u1.Function, o, &o);
                                                return true;
                                        }
                                }
                        }
                }
                return false;
    }

        static Hooker& instance() {
                static Hooker instance;
                return instance;
        }

        Hashtable<FARPROC, FARPROC> proctable;

        Hooker() {
        }
public:
        static void hook(HMODULE module, FARPROC oldproc, FARPROC newproc) {
                modifyModuleProc(module, oldproc, newproc);
        }
        static void setup(FARPROC oldproc, FARPROC newproc) {
                instance().proctable.put(oldproc, newproc);
        }
        static void hookAll(HMODULE module) {
                Pointer< Enumeration<FARPROC> > oldprocs = instance().proctable.keys();
                while (oldprocs->hasMoreElements()) {
                        FARPROC oldproc = oldprocs->nextElement();
                        FARPROC newproc = instance().proctable.get(oldproc);
                        modifyModuleProc(module, oldproc, newproc);
                }
        }
        static void unhookAll(HMODULE module) {
                Pointer< Enumeration<FARPROC> > oldprocs = instance().proctable.keys();
                while (oldprocs->hasMoreElements()) {
                        FARPROC oldproc = oldprocs->nextElement();
                        FARPROC newproc = instance().proctable.get(oldproc);
                        modifyModuleProc(module, newproc, oldproc);
                }
        }
        static FARPROC find(FARPROC original) {
                return instance().proctable.get(original);
        }
};

#endif//_YEBISOCKS_HOOKER_H_