smart vmt-hooking library for C as single header

vmt_context.h

#ifndef __VMT_CONTEXT_H__
#define __VMT_CONTEXT_H__

#ifdef _MSC_VER
#include <Windows.h>
#endif

#include <stdint.h>
#include <stdbool.h>

typedef struct VMTContext_s
{
    /// <summary> 
    /// Pointer to the original class (example: foo*) casted as two-level-pointer. 
    /// </summary>
    uintptr_t** class_base;
    /// <summary> 
    /// Original virtual function table pointer. 
    /// </summary>
    uintptr_t*  original;
    /// <summary>
    /// Our custom virtual function table pointer. 
    /// </summary>
    uintptr_t*  copy;
    /// <summary>
    /// Number of virtual functions available. 
    /// </summary>
    size_t      count;
    /// <summary> 
    /// Protection flags (optional).
    /// </summary>
    uint32_t    protection;
}VMTContext_t;

///-------------------------------------------------------------------------------------------------
/// <summary>   Checks if the data of a vmt context is valid </summary>
///
/// <remarks>   ReactiioN, 04.20.2012. </remarks>
///
/// <param name="context">  The context. </param>
///
/// <returns>   True if it succeeds, false if it fails. </returns>
///-------------------------------------------------------------------------------------------------
__forceinline
bool vmt_context_is_valid(
    const VMTContext_t* const context )
{
    return !context
        ? false
        : context->class_base && context->original && context->copy;
}

///-------------------------------------------------------------------------------------------------
/// <summary>   Changes the vtable pointer. </summary>
///
/// <remarks>   ReactiioN, 04.20.2012. </remarks>
///
/// <param name="context">  [in,out] The context. </param>
/// <param name="active">   True to active. </param>
///
/// <returns>   True if it succeeds, false if it or VirtualProtect fails. </returns>
///-------------------------------------------------------------------------------------------------
__forceinline
bool vmt_apply(
    VMTContext_t* const context,
    const bool          active )
{
    if( vmt_context_is_valid( context ) ) {
        if( context->protection != 0 ) {
            DWORD old_protect = 0;
            if( !VirtualProtect( context->class_base, sizeof( uintptr_t ), PAGE_EXECUTE_READWRITE, &old_protect ) ) {
                return false;
            }

            *context->class_base = active
                ? context->copy
                : context->original;
            
            if( !VirtualProtect( context->class_base, sizeof( uintptr_t ), old_protect, &old_protect ) ) {
                return false;
            }
        }
        else {
            *context->class_base = active
                ? context->copy
                : context->original;
        }
        return true;
    }
    return false;
}

///-------------------------------------------------------------------------------------------------
/// <summary>   Restore vmt and free memory of a vmt context </summary>
///
/// <remarks>   ReactiioN, 04.20.2012. </remarks>
///
/// <param name="context">  [in,out] The context. </param>
///-------------------------------------------------------------------------------------------------
__forceinline
void vmt_shutdown( 
    VMTContext_t* const context )
{
    if( context ) {
        vmt_apply( context, false );
        if( context->copy ) {
            free( context->copy );
        }
        context->class_base = NULL;
        context->original   = NULL;
        context->copy       = NULL;
        context->count      = 0;
    }
}

///-------------------------------------------------------------------------------------------------
/// <summary>   Initializes a vmt context based on the passed classpointer. </summary>
///
/// <remarks>   ReactiioN, 04.20.2012. </remarks>
///
/// <param name="context">          [in,out] The context. </param>
/// <param name="class_pointer">    The class pointer. </param>
/// <param name="count">            (Optional) Number of virtual functions. </param>
/// <param name="protection">       (Optional) The protection flags. </param>
///
/// <returns>   True if it succeeds, false if it fails. </returns>
///-------------------------------------------------------------------------------------------------
__forceinline
bool vmt_initialize(
    VMTContext_t* const context,
    const void* const   class_pointer,
    const size_t        count,
    const uint32_t      protection )
{
    if( !context || !class_pointer ) {
        return false;
    }

    vmt_shutdown( context );
    context->class_base = (uintptr_t**)class_pointer;
    context->original   = *context->class_base;

    context->count = count;
    if( !context->count ) {
        while( (uintptr_t*)context->original[ context->count ] ) {
        #ifdef _MSC_VER
            if( IsBadCodePtr( (FARPROC)context->original[ context->count ] ) ) {
                break;
            }
        #endif
            context->count += 1;
        }
    }    
    if( !context->count ) {
        return false;
    }

    const size_t data_size = context->count * sizeof( uintptr_t );
    context->protection    = protection;
    context->copy          = (uintptr_t*)malloc( data_size );
    memcpy( context->copy, context->original, data_size );

    return true;
}

///-------------------------------------------------------------------------------------------------
/// <summary>   Replaces a virtual function inside a vmt context. </summary>
///
/// <remarks>   ReactiioN, 04.20.2012. </remarks>
///
/// <param name="context">              [in,out] The context. </param>
/// <param name="function_index">       Zero-based index of the function. </param>
/// <param name="function_to_replace">  The function to replace. </param>
///
/// <returns>   The original address if succeeds, 0 otherwise. </returns>
///-------------------------------------------------------------------------------------------------
__forceinline
uintptr_t vmt_replace_method(
    VMTContext_t* const context,
    const size_t        function_index,
    const void* const   function_to_replace )
{
    if( vmt_context_is_valid( context ) && function_to_replace ) {
        if( function_index < context->count ) {
            context->copy[ function_index ] = (uintptr_t)function_to_replace;
            return context->original[ function_index ];
        }
    }
    return 0;
}

///-------------------------------------------------------------------------------------------------
/// <summary>   Gets the address of a virtual function </summary>
///
/// <remarks>   ReactiioN, 04.20.2012. </remarks>
///
/// <param name="context">          The context. </param>
/// <param name="function_index">   Zero-based index of the function. </param>
///
/// <returns>   The address of function at <function_index> if succeeds, 0 otherwise. </returns>
///-------------------------------------------------------------------------------------------------
__forceinline
uintptr_t vmt_get_function_address(
    const VMTContext_t* const context,
    const size_t              function_index )
{
    if( vmt_context_is_valid( context ) ) {
        if( function_index < context->count ) {
            return context->original[ function_index ];
        }
    }
    return 0;
}
#endif