#include <windows.h>
#include <tntapi.h>

#define DllImport __declspec (dllimport)

DllImport void xprintf(UCHAR *strp, ... );

typedef void (*FUNCPTR)();

LPVOID WINAPI tVirtualAlloc(LPVOID lpvAddress, DWORD cbSize,
                            DWORD fdwAllocationType, DWORD fdwProtect);

LPVOID (WINAPI *v_alloc_ptr)(LPVOID lpvAddress, DWORD cbSize,
                            DWORD fdwAllocationType, DWORD fdwProtect);

long valloc_count = 0;

BOOL WINAPI DllEntryPoint(HINSTANCE DllHandle, DWORD Reason, LPVOID reserved)
{
HMODULE dll_handle;

    switch (Reason)
    {
        case DLL_PROCESS_ATTACH :
             // get function pointer to existing Virtual Alloc()
             dll_handle = GetModuleHandle("kernel32.dll");
             if (dll_handle == NULL)
                {
                    xprintf("could not get handle of dll\n");
                    return FALSE;
                }
             v_alloc_ptr = (LPVOID) GetProcAddress(dll_handle,"VirtualAlloc");
             if (v_alloc_ptr == NULL)
                {
                    xprintf("could not get address of Virtual Alloc\n");
                    return FALSE;
                }
             // install the new (trace) function
             if (!SetProcAddress(dll_handle, "VirtualAlloc", (FUNCPTR) tVirtualAlloc))
                {
                    xprintf("error from SetProcAddress\n");
                    return FALSE;
                }
             break;
        case DLL_PROCESS_DETACH :
             xprintf("Total Number of Calls to VirtualAlloc(): %d\n",
                     valloc_count);
             break;
        case DLL_THREAD_ATTACH :  
             break;
        case DLL_THREAD_DETACH :  
             break;
        default :
             break;
    }
    return TRUE;
}

// Trace Version of Virtual Alloc 
LPVOID WINAPI tVirtualAlloc(LPVOID lpvAddress, DWORD cbSize,
                            DWORD fdwAllocationType, DWORD fdwProtect)
{
    LPVOID addr;

    // save registers
    _asm pushad

    addr = (*v_alloc_ptr) (lpvAddress, cbSize, fdwAllocationType, fdwProtect);
    valloc_count++;

    // restore registers
    _asm popad
 
    return addr;
}
