Example implementation of dxgi.dll hook to provide an imgui overlay on directx12

d3d12_imgui_hook.cpp

#define WIN32_LEAN_AND_MEAN
#define VC_EXTRA_LEAN

#define NOGDICAPMASKS         //CC_ * , LC_*, PC_*, CP_*, TC_*, RC_
//#define NOVIRTUALKEYCODES     //VK_ *
//#define NOWINMESSAGES         //WM_ * , EM_*, LB_*, CB_*
//#define NOWINSTYLES           //WS_ * , CS_*, ES_*, LBS_*, SBS_*, CBS_*
#define NOSYSMETRICS          //SM_ *
#define NOMENUS               //MF_ *
#define NOICONS               //IDI_ *
#define NOKEYSTATES           //MK_ *
#define NOSYSCOMMANDS         //SC_ *
#define NORASTEROPS           //Binary and Tertiary raster ops
#define NOSHOWWINDOW          //SW_ *
#define NOATOM                //Atom Manager routines
#define NOCLIPBOARD           //Clipboard routines
#define NOCOLOR               //Screen colors
#define NOCTLMGR              //Control and Dialog routines
#define NODRAWTEXT            //DrawText() and DT_*
#define NOGDI                 //All GDI defines and routines
#define NOKERNEL              //All KERNEL defines and routines
#define NONLS                 //All NLS defines and routines
//#define NOUSER                //All USER defines and routines
//#define NOMB                  //MB_ * and MessageBox()
#define NOMEMMGR              //GMEM_ * , LMEM_*, GHND, LHND, associated routines
#define NOMETAFILE            //typedef METAFILEPICT
#define NOMINMAX              //Macros min(a, b) and max(a, b)
//#define NOMSG                 //typedef MSG and associated routines
#define NOOPENFILE            //OpenFile(), OemToAnsi, AnsiToOem, and OF_*
#define NOSCROLL              //SB_ * and scrolling routines
#define NOSERVICE             //All Service Controller routines, SERVICE_ equates, etc.
#define NOSOUND               //Sound driver routines
#define NOTEXTMETRIC          //typedef TEXTMETRIC and associated routines
#define NOWH                  //SetWindowsHook and WH_*
//#define NOWINOFFSETS          //GWL_ * , GCL_*, associated routines
#define NOCOMM                //COMM driver routines
#define NOKANJI               //Kanji support stuff.
#define NOHELP                //Help engine interface.
#define NOPROFILER            //Profiler interface.
#define NODEFERWINDOWPOS      //DeferWindowPos routines
#define NOMCX                 //Modem Configuration Extensions

#include <Windows.h>
#include <d3d12.h>
#include <dxgi.h>
#include <dxgi1_4.h>

#include "imgui/imgui.h"
#include "imgui/imgui_impl_dx12.h"
#include "imgui/imgui_impl_win32.h"

#define EXPORT extern "C"

extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);
static LRESULT FAKE_Wndproc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam);
static LRESULT(*OLD_Wndproc)(HWND, UINT, WPARAM, LPARAM) = nullptr;

[[noreturn]] static void fatal_error(const char* const msg) {
        MessageBoxA(NULL, msg, "Fatal Error", 0);
        abort();
}

#define MAKE_STUB(NAME) void NAME(){ fatal_error(#NAME); }
MAKE_STUB(ApplyCompatResolutionQuirking);
MAKE_STUB(CompatString);
//MAKE_STUB(CompatValue);
MAKE_STUB(DXGIDumpJournal);
MAKE_STUB(PIXBeginCapture);
MAKE_STUB(PIXEndCapture);
MAKE_STUB(PIXGetCaptureState);
MAKE_STUB(SetAppCompatStringPointer);
MAKE_STUB(UpdateHMDEmulationStatus);
//MAKE_STUB(CreateDXGIFactory);
//MAKE_STUB(CreateDXGIFactory1);
//MAKE_STUB(CreateDXGIFactory2);
MAKE_STUB(DXGID3D10CreateDevice);
MAKE_STUB(DXGID3D10CreateLayeredDevice);
MAKE_STUB(DXGID3D10GetLayeredDeviceSize);
MAKE_STUB(DXGID3D10RegisterLayers);
MAKE_STUB(DXGIDeclareAdapterRemovalSupport);
MAKE_STUB(DXGIGetDebugInterface1);
MAKE_STUB(DXGIReportAdapterConfiguration);

template<typename T>
static inline T hack_vtable(T* old_fun, T new_fun) noexcept {
        T ret = *old_fun;
        DWORD perm;
        uintptr_t addr2 = (uintptr_t)old_fun;
        addr2 &= ~4095;
        VirtualProtect((void*)addr2, 4096, PAGE_EXECUTE_READWRITE, &perm);
        *old_fun = new_fun;
        VirtualProtect((void*)addr2, 4096, perm, nullptr);
        return ret;
}

template<typename T>
static inline void GetDXGIProc(const char* proc_name, T* func_out) noexcept {
        if (*func_out != nullptr) return;

        static HMODULE dxgi_handle = nullptr;

        if (dxgi_handle == nullptr) {
                dxgi_handle = LoadLibraryA("c:\\Windows\\System32\\dxgi.dll");
                assert(dxgi_handle != nullptr);
        }
        *func_out = (T) GetProcAddress(dxgi_handle, proc_name);
}

EXPORT HRESULT FAKE_CompatValue() {
        return 0;
}

typedef HRESULT (*FUN_CreateSwapChainForHwnd)(
        IDXGIFactory4 *This,
        IUnknown *pDevice,
        HWND hWnd,
        const DXGI_SWAP_CHAIN_DESC1 *pDesc,
        const DXGI_SWAP_CHAIN_FULLSCREEN_DESC *pFullscreenDesc,
        IDXGIOutput *pRestrictToOutput,
        IDXGISwapChain1 **ppSwapChain);

static HRESULT FAKE_CreateSwapChainForHwnd(
        IDXGIFactory4 *This,
        IUnknown *pDevice,
        HWND hWnd,
        const DXGI_SWAP_CHAIN_DESC1 *pDesc,
        const DXGI_SWAP_CHAIN_FULLSCREEN_DESC *pFullscreenDesc,
        IDXGIOutput *pRestrictToOutput,
        IDXGISwapChain1 **ppSwapChain);

static FUN_CreateSwapChainForHwnd OLD_CreateSwapChainForHwnd = nullptr;


typedef HRESULT (*FUN_CreateSwapChain)(IDXGIFactory * This, IUnknown *pDevice, DXGI_SWAP_CHAIN_DESC *pDesc, IDXGISwapChain **ppSwapChain);
static HRESULT FAKE_CreateSwapChain(IDXGIFactory * This, IUnknown *pDevice, DXGI_SWAP_CHAIN_DESC *pDesc, IDXGISwapChain **ppSwapChain);
static FUN_CreateSwapChain OLD_CreateSwapChain = nullptr;
typedef void(*dummy_function)(void);
#define DEBUG_VTABLE(NAME) hack_vtable(&fake->vTable->NAME, (dummy_function)[](void)->void{fatal_error(#NAME); });

static void hack_factory(void **ppFactory) {
        struct {
                struct {
                        dummy_function QueryInterface;
                        dummy_function AddRef;
                        dummy_function Release;
                        dummy_function SetPrivateData;
                        dummy_function SetPrivateDataInterface;
                        dummy_function GetPrivateData;
                        dummy_function GetParent;
                        dummy_function EnumAdapters;
                        dummy_function MakeWindowAssociation;
                        dummy_function GetWindowAssociation;
                        FUN_CreateSwapChain CreateSwapChain;
                        dummy_function CreateSoftwareAdapter;
                        dummy_function EnumAdapters1;
                        dummy_function IsCurrent;
                        dummy_function IsWindowedStereoEnabled;
                        FUN_CreateSwapChainForHwnd CreateSwapChainForHwnd;
                        dummy_function CreateSwapChainForCoreWindow;
                        dummy_function GetSharedResourceAdapterLuid;
                        dummy_function RegisterStereoStatusWindow;
                        dummy_function RegisterStereoStatusEvent;
                        dummy_function UnregisterStereoStatus;
                        dummy_function RegisterOcclusionStatusWindow;
                        dummy_function RegisterOcclusionStatusEvent;
                        dummy_function UnregisterOcclusionStatus;
                        dummy_function CreateSwapChainForComposition;
                        dummy_function EnumAdapterByLuid;
                        dummy_function EnumWarpAdapter;
                } *vTable;
        } *fake = *(decltype(fake)*)ppFactory;

        if (!OLD_CreateSwapChainForHwnd) {
                OLD_CreateSwapChainForHwnd = hack_vtable(&fake->vTable->CreateSwapChainForHwnd, FAKE_CreateSwapChainForHwnd);
                OLD_CreateSwapChain = hack_vtable(&fake->vTable->CreateSwapChain, FAKE_CreateSwapChain);
                DEBUG_VTABLE(CreateSwapChainForCoreWindow);
                DEBUG_VTABLE(CreateSwapChainForComposition);
        }
}

EXPORT HRESULT FAKE_CreateDXGIFactory(REFIID riid, void **ppFactory) {
        static decltype(FAKE_CreateDXGIFactory)* func = nullptr;
        GetDXGIProc("CreateDXGIFactory", &func);
        auto ret = func(riid, ppFactory);
        assert(ret == S_OK);
        hack_factory(ppFactory);
        return ret;
}


EXPORT HRESULT FAKE_CreateDXGIFactory1(REFIID riid, void **ppFactory) {
        static decltype(FAKE_CreateDXGIFactory1)* func = nullptr;
        GetDXGIProc("CreateDXGIFactory1", &func);
        auto ret = func(riid, ppFactory);
        assert(ret == S_OK);
        hack_factory(ppFactory);
        return ret;
}

EXPORT HRESULT FAKE_CreateDXGIFactory2(UINT Flags, REFIID riid, void **ppFactory) {
        static decltype(FAKE_CreateDXGIFactory2)* func = nullptr;
        GetDXGIProc("CreateDXGIFactory2", &func);
        auto ret = func(Flags, riid, ppFactory);
        assert(ret == S_OK);
        hack_factory(ppFactory);
        return ret;
}

EXPORT BOOL WINAPI DllMain(HINSTANCE, DWORD, LPVOID) {
        return TRUE;
}

typedef HRESULT(*FUN_Present)(IDXGISwapChain3 * This, UINT SyncInterval, UINT PresentFlags);
static HRESULT FAKE_Present(IDXGISwapChain3 * This, UINT SyncInterval, UINT PresentFlags);
static FUN_Present OLD_Present = nullptr;
static ID3D12CommandQueue* d3d12CommandQueue = nullptr;


static void hack_swapchain(void ** ppSwapChain) {
        struct {
                struct {
                        dummy_function QueryInterface;
                        dummy_function AddRef;
                        dummy_function Release;
                        dummy_function SetPrivateData;
                        dummy_function SetPrivateDataInterface;
                        dummy_function GetPrivateData;
                        dummy_function GetParent;
                        dummy_function GetDevice;
                        FUN_Present Present;
                        dummy_function GetBuffer;
                        dummy_function SetFullscreenState;
                        dummy_function GetFullscreenState;
                        dummy_function GetDesc;
                        dummy_function ResizeBuffers;
                        dummy_function ResizeTarget;
                        dummy_function GetContainingOutput;
                        dummy_function GetFrameStatistics;
                        dummy_function GetLastPresentCount;
                        dummy_function GetDesc1;
                        dummy_function GetFullscreenDesc;
                        dummy_function GetHwnd;
                        dummy_function Present1;
                        dummy_function IsTemporaryMonoSupported;
                        dummy_function GetRestrictToOutput;
                        dummy_function SetBackgroundColor;
                        dummy_function GetBackgroundColor;
                        dummy_function SetRotation;
                        dummy_function GetRotation;
                        dummy_function SetSourceSize;
                        dummy_function GetSourceSize;
                        dummy_function SetMaximumFrameLatency;
                        dummy_function GetMaximumFrameLatency;
                        dummy_function GetFrameLatencyWaitableObject;
                        dummy_function SetMatrixTransform;
                        dummy_function GetMatrixTransform;
                        dummy_function GetCurrentBackBufferIndex;
                        dummy_function CheckColorSpaceSupport;
                        dummy_function SetColorSpace1;
                        dummy_function ResizeBuffers1;
                } *vTable;
        } *fake = *(decltype(fake)*)ppSwapChain;

        if (!OLD_Present) {
                OLD_Present = hack_vtable(&fake->vTable->Present, FAKE_Present);
                DEBUG_VTABLE(Present1);
        }
}


static HRESULT FAKE_CreateSwapChain(IDXGIFactory * This, IUnknown *pDevice, DXGI_SWAP_CHAIN_DESC *pDesc, IDXGISwapChain **ppSwapChain) {
        d3d12CommandQueue = (ID3D12CommandQueue*)pDevice;
        auto ret = OLD_CreateSwapChain(This, pDevice, pDesc, ppSwapChain);
        hack_swapchain((void**)ppSwapChain);
        return ret;
}


static HRESULT FAKE_CreateSwapChainForHwnd(
        IDXGIFactory4 *This,
        IUnknown *pDevice,
        HWND hWnd,
        const DXGI_SWAP_CHAIN_DESC1 *pDesc,
        const DXGI_SWAP_CHAIN_FULLSCREEN_DESC *pFullscreenDesc,
        IDXGIOutput *pRestrictToOutput,
        IDXGISwapChain1 **ppSwapChain)
{
        d3d12CommandQueue = (ID3D12CommandQueue*)pDevice;
        auto ret = OLD_CreateSwapChainForHwnd(This, pDevice, hWnd, pDesc, pFullscreenDesc, pRestrictToOutput, ppSwapChain);
        hack_swapchain((void**)ppSwapChain);
        return ret;
}

static bool should_show_ui = true;

static HRESULT FAKE_Present(IDXGISwapChain3 * This, UINT SyncInterval, UINT PresentFlags) {
        struct FrameContext {
                ID3D12CommandAllocator* commandAllocator = nullptr;
                ID3D12Resource* main_render_target_resource = nullptr;
                D3D12_CPU_DESCRIPTOR_HANDLE main_render_target_descriptor;
        };

        static ID3D12Device* d3d12Device = nullptr;
        static ID3D12DescriptorHeap* d3d12DescriptorHeapBackBuffers = nullptr;
        static ID3D12DescriptorHeap* d3d12DescriptorHeapImGuiRender = nullptr;
        static ID3D12GraphicsCommandList* d3d12CommandList = nullptr;
        static HWND window_handle = nullptr;
        static UINT buffersCounts = 0;
        static FrameContext* frameContext = nullptr;


        if (!d3d12Device) {
                IMGUI_CHECKVERSION();
                ImGui::CreateContext();
                ImGuiIO& io = ImGui::GetIO();
                io.MouseDrawCursor = true;
                ImGui::StyleColorsDark();

                CreateEvent(nullptr, false, false, nullptr);

                DXGI_SWAP_CHAIN_DESC sdesc;
                This->GetDevice(__uuidof(ID3D12Device), (void**)&d3d12Device);
                This->GetDesc(&sdesc);
                This->GetHwnd(&window_handle);
                sdesc.Flags = DXGI_SWAP_CHAIN_FLAG_ALLOW_MODE_SWITCH;
                sdesc.OutputWindow = window_handle;
                sdesc.Windowed = ((GetWindowLongPtr(window_handle, GWL_STYLE) & WS_POPUP) != 0) ? false : true;

                buffersCounts = sdesc.BufferCount;
                frameContext = (FrameContext*) calloc(buffersCounts, sizeof(FrameContext));

                D3D12_DESCRIPTOR_HEAP_DESC descriptorImGuiRender = {};
                descriptorImGuiRender.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
                descriptorImGuiRender.NumDescriptors = buffersCounts;
                descriptorImGuiRender.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;

                d3d12Device->CreateDescriptorHeap(&descriptorImGuiRender, IID_PPV_ARGS(&d3d12DescriptorHeapImGuiRender));

                ID3D12CommandAllocator* allocator;
                d3d12Device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, IID_PPV_ARGS(&allocator));
                for (size_t i = 0; i < buffersCounts; i++) {
                        frameContext[i].commandAllocator = allocator;
                }

                d3d12Device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, allocator, NULL, IID_PPV_ARGS(&d3d12CommandList));

                D3D12_DESCRIPTOR_HEAP_DESC descriptorBackBuffers;
                descriptorBackBuffers.Type = D3D12_DESCRIPTOR_HEAP_TYPE_RTV;
                descriptorBackBuffers.NumDescriptors = buffersCounts;
                descriptorBackBuffers.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
                descriptorBackBuffers.NodeMask = 1;

                d3d12Device->CreateDescriptorHeap(&descriptorBackBuffers, IID_PPV_ARGS(&d3d12DescriptorHeapBackBuffers));

                const auto rtvDescriptorSize = d3d12Device->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
                D3D12_CPU_DESCRIPTOR_HANDLE rtvHandle = d3d12DescriptorHeapBackBuffers->GetCPUDescriptorHandleForHeapStart();

                for (UINT i = 0; i < buffersCounts; i++) {
                        ID3D12Resource* pBackBuffer = nullptr;
                        frameContext[i].main_render_target_descriptor = rtvHandle;
                        This->GetBuffer(i, IID_PPV_ARGS(&pBackBuffer));
                        d3d12Device->CreateRenderTargetView(pBackBuffer, nullptr, rtvHandle);
                        frameContext[i].main_render_target_resource = pBackBuffer;
                        rtvHandle.ptr += rtvDescriptorSize;
                }

                ImGui_ImplWin32_Init(window_handle);
                ImGui_ImplDX12_Init(d3d12Device, buffersCounts,
                        DXGI_FORMAT_R8G8B8A8_UNORM, d3d12DescriptorHeapImGuiRender,
                        d3d12DescriptorHeapImGuiRender->GetCPUDescriptorHandleForHeapStart(),
                        d3d12DescriptorHeapImGuiRender->GetGPUDescriptorHandleForHeapStart());

                ImGui_ImplDX12_CreateDeviceObjects();
                OLD_Wndproc = (decltype(OLD_Wndproc)) SetWindowLongPtrW(window_handle, GWLP_WNDPROC, (LONG_PTR) FAKE_Wndproc);
        }

        if (should_show_ui && d3d12CommandQueue) {
                ImGui_ImplDX12_NewFrame();
                ImGui_ImplWin32_NewFrame();
                ImGui::NewFrame();

                ImGui::ShowDemoWindow();

                FrameContext& currentFrameContext = frameContext[This->GetCurrentBackBufferIndex()];
                currentFrameContext.commandAllocator->Reset();

                D3D12_RESOURCE_BARRIER barrier;
                barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
                barrier.Flags = D3D12_RESOURCE_BARRIER_FLAG_NONE;
                barrier.Transition.pResource = currentFrameContext.main_render_target_resource;
                barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
                barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_PRESENT;
                barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_RENDER_TARGET;

                
                d3d12CommandList->Reset(currentFrameContext.commandAllocator, nullptr);
                d3d12CommandList->ResourceBarrier(1, &barrier);
                d3d12CommandList->OMSetRenderTargets(1, &currentFrameContext.main_render_target_descriptor, FALSE, nullptr);
                d3d12CommandList->SetDescriptorHeaps(1, &d3d12DescriptorHeapImGuiRender);

                ImGui::Render();
                ImGui_ImplDX12_RenderDrawData(ImGui::GetDrawData(), d3d12CommandList);

                barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_RENDER_TARGET;
                barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_PRESENT;

                d3d12CommandList->ResourceBarrier(1, &barrier);
                d3d12CommandList->Close();

                d3d12CommandQueue->ExecuteCommandLists(1, reinterpret_cast<ID3D12CommandList* const*>(&d3d12CommandList));
        }

        return OLD_Present(This, SyncInterval, PresentFlags);
}

static LRESULT FAKE_Wndproc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
        if ((uMsg == WM_KEYUP) && (wParam == VK_ESCAPE)) {
                should_show_ui = !should_show_ui;
                return true;
        }

        if (should_show_ui && ImGui_ImplWin32_WndProcHandler(hWnd, uMsg, wParam, lParam)) {
                return true;
        }

        return CallWindowProc((WNDPROC)OLD_Wndproc, hWnd, uMsg, wParam, lParam);
}