RemoteExec.h

#pragma once

#include "../../Include/Winheaders.h"
#include "../../Asm/AsmHelper.h"
#include "../Threads/Threads.h"
#include "../MemBlock.h"


// User data offsets
#define INTRET_OFFSET   0x00
#define RET_OFFSET      0x08
#define ERR_OFFSET      0x10
#define EVENT_OFFSET    0x18
#define ARGS_OFFSET     0x20


namespace blackbone
{

class RemoteExec
{
    template<typename Fn>
    friend class RemoteFuncBase;

    typedef std::vector<AsmVariant> vecArgs;

public:
    BLACKBONE_API RemoteExec( class Process& proc );
    BLACKBONE_API ~RemoteExec();

    BLACKBONE_API NTSTATUS CreateRPCEnvironment( bool bThread = true, bool bEvent = true );

    BLACKBONE_API NTSTATUS ExecInNewThread( PVOID pCode, size_t size, uint64_t& callResult );

    BLACKBONE_API NTSTATUS ExecInWorkerThread( PVOID pCode, size_t size, uint64_t& callResult );

    BLACKBONE_API NTSTATUS ExecInAnyThread( PVOID pCode, size_t size, uint64_t& callResult, Thread& thread );

    BLACKBONE_API DWORD ExecDirect( ptr_t pCode, ptr_t arg );

    BLACKBONE_API void AddReturnWithEvent(
        AsmHelperBase& a,
        eModType mt = mt_default, 
        eReturnType retType = rt_int32,
        uint32_t retOffset = RET_OFFSET 
        );

    BLACKBONE_API inline NTSTATUS GetLastStatus()
    {
        return _userData.Read<NTSTATUS>( ERR_OFFSET, STATUS_NOT_FOUND );
    }

    BLACKBONE_API void TerminateWorker();

    BLACKBONE_API inline Thread* getWorker() { return &_hWorkThd; }

    BLACKBONE_API inline class ProcessMemory& memory() { return _memory; }

    BLACKBONE_API void reset();

private:

    DWORD CreateWorkerThread();

    bool CreateAPCEvent( DWORD threadID );

    NTSTATUS CopyCode( PVOID pCode, size_t size );

    BLACKBONE_API bool PrepareCallAssembly( AsmHelperBase& a, const void* pfn,
                                            std::vector<blackbone::AsmVariant>& args,
                                            eCalligConvention cc, eReturnType retType );

#pragma warning(disable : 4127)

    template<typename T>
    inline bool GetCallResult( T& result )
    { 
        if (sizeof(T) > sizeof(uint64_t))
        {
            if (std::is_reference<T>::value)
                return _userData.Read( _userData.Read<size_t>( RET_OFFSET, 0 ), sizeof(T), (PVOID)&result ) == STATUS_SUCCESS;
            else
                return _userData.Read( ARGS_OFFSET, sizeof(T), (PVOID)&result ) == STATUS_SUCCESS;
        }
        else
            return _userData.Read( RET_OFFSET, sizeof(T), (PVOID)&result ) == STATUS_SUCCESS;
    }
#pragma warning(default : 4127)

    RemoteExec( const RemoteExec& ) = delete;
    RemoteExec& operator =(const RemoteExec&) = delete;

private:    
    // Process routines
    class Process&        _proc;
    class ProcessModules& _mods;
    class ProcessMemory&  _memory;
    class ProcessThreads& _threads;

    Thread   _hWorkThd;         // Worker thread handle
    HANDLE   _hWaitEvent;       // APC sync event handle
    MemBlock _workerCode;       // Worker thread address space
    MemBlock _userCode;         // Codecave for code execution
    MemBlock _userData;         // Region to store copied structures and strings
    bool     _apcPatched;       // KiUserApcDispatcher was patched
};


}