TraceHook.h

#pragma once

#include "../Include/WinHeaders.h"

#include <stdint.h>
#include <vector>
#include <map>
#include <unordered_map>

namespace blackbone
{

enum TraceState
{
    TS_Start,       // Initial state. Internal use only
    TS_Step,        // Do single-step
    TS_StepOut,     // Break on function return
    TS_StepInto,    // Step into specific function
    TS_WaitReturn,  // Wait for break-on-return
};

struct PathNode
{
    TraceState action;
    uintptr_t arg;

    PathNode( TraceState  _action, uintptr_t _arg = 0 )
        : action( _action )
        , arg( _arg ) { }
};


struct HookContext
{
    typedef std::unordered_map<uintptr_t, std::pair<uintptr_t, bool>> mapHooks;
    typedef std::vector<PathNode> vecState;

    uintptr_t lastIP = 0;       // Previous EIP/RIP value
    uintptr_t lastSP = 0;       // Previous ESP/RSP value
    uintptr_t targetPtr = 0;    // Address causing exception
    uintptr_t origPtrVal = 0;   // Original pointer value
    uintptr_t checkIP = 0;      // Address of instruction that checks target pointer
    uintptr_t breakValue = 0;   // Value used to generate exception
    uintptr_t stateIdx = 0;     // Current state index in state vector

    TraceState state = TS_Start;    // Current tracing state
    vecState tracePath;             // Function trace path
    mapHooks hooks;                 // List of hooks associated with current pointer


    void reset()
    {
        state = TS_Start;
        lastIP = lastSP = 0;
        stateIdx = 0;

        // Mark hooks as non-called
        for (auto& item : hooks)
            item.second.second = false;
    }
};

class TraceHook
{

public:
    typedef std::map<uintptr_t, HookContext> mapContext;
    typedef std::vector <std::pair<uintptr_t, uintptr_t>> vecStackFrames;

public:
    ~TraceHook();
    static TraceHook& Instance();

    BLACKBONE_API bool ApplyHook( void* targetFunc,
                                  void* hookFunc,
                                  void* ptrAddress,
                                  const HookContext::vecState& tracePath = HookContext::vecState(),
                                  void* checkIP = 0 );

    BLACKBONE_API bool RemoveHook( void* targetFunc );

private:
    // 
    // Singleton
    //
    TraceHook();
    TraceHook( const TraceHook& ) = delete;
    TraceHook& operator =( const TraceHook& ) = delete;

    //
    // Exception handlers
    //
    static LONG __stdcall VecHandler( PEXCEPTION_POINTERS ExceptionInfo );
    LONG VecHandlerP( PEXCEPTION_POINTERS ExceptionInfo );

    size_t StackBacktrace( uintptr_t ip, uintptr_t sp, vecStackFrames& results, uintptr_t depth = 10 );

    inline void BreakOnReturn( uintptr_t sp );

    bool CheckBranching( const HookContext& ctx, uintptr_t ip, uintptr_t sp );

    void HandleBranch( HookContext& ctx, PCONTEXT exptContex );

    bool RestorePtr( const HookContext& ctx, PEXCEPTION_POINTERS ExceptionInfo );

private:
    PVOID       _pExptHandler = nullptr;        // Exception handler
    mapContext  _contexts;                      // Hook contexts
    uintptr_t   _breakPtr = 0x2000;             // Exception pointer generator
};

}