mhook-test.cpp

//Copyright (c) 2007-2008, Marton Anka
//
//Permission is hereby granted, free of charge, to any person obtaining a 
//copy of this software and associated documentation files (the "Software"), 
//to deal in the Software without restriction, including without limitation 
//the rights to use, copy, modify, merge, publish, distribute, sublicense, 
//and/or sell copies of the Software, and to permit persons to whom the 
//Software is furnished to do so, subject to the following conditions:
//
//The above copyright notice and this permission notice shall be included 
//in all copies or substantial portions of the Software.
//
//THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 
//OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
//FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 
//THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
//LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
//FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 
//IN THE SOFTWARE.

#include "stdafx.h"
#include "mhook-lib/mhook.h"

//=========================================================================
// Define _NtOpenProcess so we can dynamically bind to the function
//
typedef struct _CLIENT_ID 
{
    DWORD_PTR UniqueProcess;
    DWORD_PTR UniqueThread;
} CLIENT_ID, *PCLIENT_ID;

typedef ULONG (WINAPI* _NtOpenProcess)(OUT PHANDLE ProcessHandle, 
         IN ACCESS_MASK AccessMask, IN PVOID ObjectAttributes, 
         IN PCLIENT_ID ClientId ); 

//=========================================================================
// Define _SelectObject so we can dynamically bind to the function
typedef HGDIOBJ (WINAPI* _SelectObject)(HDC hdc, HGDIOBJ hgdiobj); 

//=========================================================================
// Define _getaddrinfo so we can dynamically bind to the function
typedef int (WSAAPI* _getaddrinfo)(const char* nodename, const char* servname, const struct addrinfo* hints, struct addrinfo** res);

//=========================================================================
// Define _HeapAlloc so we can dynamically bind to the function
typedef LPVOID (WINAPI *_HeapAlloc)(HANDLE, DWORD, SIZE_T);

//=========================================================================
// Define _NtClose so we can dynamically bind to the function
typedef ULONG (WINAPI* _NtClose)(IN HANDLE Handle);

//=========================================================================
// Get the current (original) address to the functions to be hooked
//
_NtOpenProcess TrueNtOpenProcess = (_NtOpenProcess)
    GetProcAddress(GetModuleHandle(L"ntdll"), "NtOpenProcess");

_SelectObject TrueSelectObject = (_SelectObject)
    GetProcAddress(GetModuleHandle(L"gdi32"), "SelectObject");

_getaddrinfo Truegetaddrinfo = (_getaddrinfo)GetProcAddress(GetModuleHandle(L"ws2_32"), "getaddrinfo");

_HeapAlloc TrueHeapAlloc = (_HeapAlloc)GetProcAddress(GetModuleHandle(L"kernel32"), "HeapAlloc");

_NtClose TrueNtClose = (_NtClose)GetProcAddress(GetModuleHandle(L"ntdll"), "NtClose");

//=========================================================================
// This is the function that will replace NtOpenProcess once the hook 
// is in place
//
ULONG WINAPI HookNtOpenProcess(OUT PHANDLE ProcessHandle, 
                               IN ACCESS_MASK AccessMask, 
                               IN PVOID ObjectAttributes, 
                               IN PCLIENT_ID ClientId)
{
    printf("***** Call to open process %d\n", ClientId->UniqueProcess);
    return TrueNtOpenProcess(ProcessHandle, AccessMask, 
        ObjectAttributes, ClientId);
}

//=========================================================================
// This is the function that will replace SelectObject once the hook 
// is in place
//
HGDIOBJ WINAPI HookSelectobject(HDC hdc, HGDIOBJ hgdiobj)
{
    printf("***** Call to SelectObject(0x%p, 0x%p)\n", hdc, hgdiobj);
    return TrueSelectObject(hdc, hgdiobj);
}

//=========================================================================
// This is the function that will replace SelectObject once the hook 
// is in place
//
int WSAAPI Hookgetaddrinfo(const char* nodename, const char* servname, const struct addrinfo* hints, struct addrinfo** res)
{
    printf("***** Call to getaddrinfo(0x%p, 0x%p, 0x%p, 0x%p)\n", nodename, servname, hints, res);
    return Truegetaddrinfo(nodename, servname, hints, res);
}

//=========================================================================
// This is the function that will replace HeapAlloc once the hook 
// is in place
//
LPVOID WINAPI HookHeapAlloc(HANDLE a_Handle, DWORD a_Bla, SIZE_T a_Bla2) 
{
    printf("***** Call to HeapAlloc(0x%p, %u, 0x%Id)\n", a_Handle, a_Bla, a_Bla2);
    return TrueHeapAlloc(a_Handle, a_Bla, a_Bla2);
}

//=========================================================================
// This is the function that will replace NtClose once the hook 
// is in place
//
ULONG WINAPI HookNtClose(HANDLE hHandle) 
{
    printf("***** Call to NtClose(0x%p)\n", hHandle);
    return TrueNtClose(hHandle);
}

//=========================================================================
// This is where the work gets done.
//
int wmain(int argc, WCHAR* argv[])
{
    HANDLE hProc = NULL;

    // Set the hook
    if (Mhook_SetHook((PVOID*)&TrueNtOpenProcess, HookNtOpenProcess)) 
    {
        // Now call OpenProcess and observe NtOpenProcess being redirected
        // under the hood.
        hProc = OpenProcess(PROCESS_ALL_ACCESS, 
            FALSE, GetCurrentProcessId());
        if (hProc) 
        {
            printf("Successfully opened self: %p\n", hProc);
            CloseHandle(hProc);
        } 
        else 
        {
            printf("Could not open self: %d\n", GetLastError());
        }
        // Remove the hook
        Mhook_Unhook((PVOID*)&TrueNtOpenProcess);
    }

    // Call OpenProces again - this time there won't be a redirection as
    // the hook has bee removed.
    hProc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, GetCurrentProcessId());
    if (hProc) 
    {
        printf("Successfully opened self: %p\n", hProc);
        CloseHandle(hProc);
    } 
    else 
    {
        printf("Could not open self: %d\n", GetLastError());
    }

    // Test another hook, this time in SelectObject
    // (SelectObject is interesting in that on XP x64, the second instruction
    // in the trampoline uses IP-relative addressing and we need to do some
    // extra work under the hood to make things work properly. This really
    // is more of a test case rather than a demo.)
    printf("Testing SelectObject.\n");
    if (Mhook_SetHook((PVOID*)&TrueSelectObject, HookSelectobject)) 
    {
        // error checking omitted for brevity. doesn't matter much 
        // in this context anyway.
        HDC hdc = GetDC(NULL);
        HDC hdcMem = CreateCompatibleDC(hdc);
        HBITMAP hbm = CreateCompatibleBitmap(hdc, 32, 32);
        HBITMAP hbmOld = (HBITMAP)SelectObject(hdcMem, hbm);
        SelectObject(hdcMem, hbmOld);
        DeleteObject(hbm);
        DeleteDC(hdcMem);
        ReleaseDC(NULL, hdc);
        // Remove the hook
        Mhook_Unhook((PVOID*)&TrueSelectObject);
    }

    printf("Testing getaddrinfo.\n");
    if (Mhook_SetHook((PVOID*)&Truegetaddrinfo, Hookgetaddrinfo)) 
    {
        // error checking omitted for brevity. doesn't matter much 
        // in this context anyway.
        WSADATA wd = {0};
        WSAStartup(MAKEWORD(2, 2), &wd);
        char* ip = "localhost";
        struct addrinfo aiHints;
        struct addrinfo *res = NULL;
        memset(&aiHints, 0, sizeof(aiHints));
        aiHints.ai_family = PF_UNSPEC;
        aiHints.ai_socktype = SOCK_STREAM;
        if (getaddrinfo(ip, NULL, &aiHints, &res)) 
        {
            printf("getaddrinfo failed\n");
        }
        else 
        {
            int n = 0;
            while(res) 
            {
                res = res->ai_next;
                n++;
            }
            printf("got %d addresses\n", n);
        }
        WSACleanup();
        // Remove the hook
        Mhook_Unhook((PVOID*)&Truegetaddrinfo);
    }

    printf("Testing HeapAlloc.\n");
    if (Mhook_SetHook((PVOID*)&TrueHeapAlloc, HookHeapAlloc))
    {
        free(malloc(10));
        // Remove the hook
        Mhook_Unhook((PVOID*)&TrueHeapAlloc);
    }

    printf("Testing NtClose.\n");
    if (Mhook_SetHook((PVOID*)&TrueNtClose, HookNtClose))
    {
        CloseHandle(NULL);
        // Remove the hook
        Mhook_Unhook((PVOID*)&TrueNtClose);
    }

    return 0;
}