overlay.cpp

// =====================================================================
// NFSU2 LiveTextures - D3D9 device hooks and the ImGui overlay
//
// How we get at the device: SPEED2.EXE imports Direct3DCreate9 by name, so
// we redirect that one entry in the game's own import table. From the
// IDirect3D9 it returns we catch CreateDevice, and the moment the device
// exists we hook EndScene/Reset (MinHook) and hand the device to
// textures.cpp. Doing it this early means we are in place before any other
// plugin wraps or re-hooks the device.
//
// What we deliberately do not do: patch exports of d3d9.dll, dinput8.dll or
// user32.dll. Those are shared by every mod in the process (ReShade, other
// ASI plugins, widescreen fixes), and two plugins detouring the same export
// is a classic way to crash on startup. The game's import table and the
// vtables of objects created for the game belong to the game alone.
//
// Input: while the overlay is open, the game's DirectInput keyboard and
// mouse read as "nothing pressed" and the window messages ImGui used are
// swallowed. With the overlay closed the EndScene hook does its per-frame
// housekeeping and returns.
// =====================================================================

#include "plugin.h"
#include <d3d9.h>
#include <dinput.h>
#include <cstdio>
#include <cstdarg>
#include <commdlg.h>
#include <shellapi.h>
#include <vector>
#include <set>
#include <algorithm>
#include "MinHook.h"
#include "imgui.h"
#include "backends/imgui_impl_dx9.h"
#include "backends/imgui_impl_win32.h"

extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND, UINT, WPARAM, LPARAM);

namespace {

typedef IDirect3D9* (WINAPI* Direct3DCreate9_t)(UINT);
typedef HRESULT(APIENTRY* CreateDevice_t)(IDirect3D9*, UINT, D3DDEVTYPE, HWND, DWORD, D3DPRESENT_PARAMETERS*, IDirect3DDevice9**);
typedef HRESULT(APIENTRY* EndScene_t)(IDirect3DDevice9*);
typedef HRESULT(APIENTRY* Reset_t)(IDirect3DDevice9*, D3DPRESENT_PARAMETERS*);
typedef HRESULT(WINAPI* DI8Create_t)(HINSTANCE, DWORD, REFIID, LPVOID*, LPUNKNOWN);
typedef HRESULT(STDMETHODCALLTYPE* DICreateDevice_t)(void*, REFGUID, void**, LPUNKNOWN);
typedef HRESULT(STDMETHODCALLTYPE* DIGetDeviceState_t)(void*, DWORD, LPVOID);
typedef HRESULT(STDMETHODCALLTYPE* DIGetDeviceData_t)(void*, DWORD, void*, LPDWORD, DWORD);

Direct3DCreate9_t  g_origDirect3DCreate9 = nullptr;
CreateDevice_t     g_origCreateDevice    = nullptr;
EndScene_t         g_origEndScene        = nullptr;
Reset_t            g_origReset           = nullptr;
DI8Create_t        g_origDI8Create       = nullptr;
DICreateDevice_t   g_origDICreateDevice  = nullptr;
DIGetDeviceState_t g_origGetDeviceState  = nullptr;
DIGetDeviceData_t  g_origGetDeviceData   = nullptr;
bool       g_deviceHooked = false;
WNDPROC    g_origWndProc  = nullptr;
HWND       g_hWnd         = nullptr;
bool       g_imguiReady   = false;
bool       g_visible      = false;
bool       g_keyWasDown   = false;
float      g_fps          = 0.f;
DWORD      g_lastTick     = 0;
int        g_frames       = 0;
char       g_status[128]  = "";
void*      g_blockedDevices[8];   // the game's DirectInput keyboard and mouse devices
int        g_blockedCount = 0;
bool       g_dikKeyDown   = false;   // toggle key as seen through the game's own DirectInput keyboard
bool       g_dikSeen      = false;   // we have seen it there at least once, so that is where we keep looking

bool OverlayOwnsInput() { return g_visible && g_imguiReady; }

bool IsBlockedDevice(void* dev)
{
    for (int i = 0; i < g_blockedCount; ++i) if (g_blockedDevices[i] == dev) return true;
    return false;
}

// Replaces one slot in an object's vtable and returns what was there. The vtable is shared, so every
// object of that class is affected - which is what we want for the DirectInput devices.
void* SwapVtableEntry(void* object, int index, void* replacement)
{
    void** slot = *reinterpret_cast<void***>(object) + index;
    DWORD old;
    if (!VirtualProtect(slot, sizeof(void*), PAGE_READWRITE, &old)) return nullptr;
    void* previous = *slot;
    *slot = replacement;
    VirtualProtect(slot, sizeof(void*), old, &old);
    return previous;
}

// Redirects one named import of SPEED2.EXE and returns the previous target. Plain PE walking: import
// directory -> descriptor of the DLL -> the name table and the address table run in parallel.
void* PatchImport(const char* dll, const char* func, void* replacement)
{
    BYTE* base = reinterpret_cast<BYTE*>(GetModuleHandleA(nullptr));
    IMAGE_NT_HEADERS* nt = reinterpret_cast<IMAGE_NT_HEADERS*>(base + reinterpret_cast<IMAGE_DOS_HEADER*>(base)->e_lfanew);
    IMAGE_DATA_DIRECTORY& dir = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (!dir.VirtualAddress) return nullptr;
    for (IMAGE_IMPORT_DESCRIPTOR* imp = reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>(base + dir.VirtualAddress); imp->Name; ++imp) {
        if (_stricmp(reinterpret_cast<const char*>(base + imp->Name), dll) != 0) continue;
        IMAGE_THUNK_DATA* names = reinterpret_cast<IMAGE_THUNK_DATA*>(base + imp->OriginalFirstThunk);
        IMAGE_THUNK_DATA* funcs = reinterpret_cast<IMAGE_THUNK_DATA*>(base + imp->FirstThunk);
        for (; names->u1.AddressOfData; ++names, ++funcs) {
            if (names->u1.Ordinal & IMAGE_ORDINAL_FLAG) continue;
            IMAGE_IMPORT_BY_NAME* byName = reinterpret_cast<IMAGE_IMPORT_BY_NAME*>(base + names->u1.AddressOfData);
            if (strcmp(byName->Name, func) != 0) continue;
            DWORD old;
            if (!VirtualProtect(&funcs->u1.Function, sizeof(DWORD), PAGE_READWRITE, &old)) return nullptr;
            void* previous = reinterpret_cast<void*>(funcs->u1.Function);
            funcs->u1.Function = reinterpret_cast<DWORD>(replacement);
            VirtualProtect(&funcs->u1.Function, sizeof(DWORD), old, &old);
            return previous;
        }
    }
    return nullptr;
}

// --- DirectInput (input blocking + overlay key) -------------------------
// The INI takes a virtual-key code, DirectInput reports scan codes. Only keys that make sense as an
// overlay toggle are mapped; for anything else we fall back to GetAsyncKeyState.
int VkToDik(int vk)
{
    switch (vk) {
    case VK_INSERT: return 0xD2; case VK_END: return 0xCF; case VK_HOME: return 0xC7; case VK_DELETE: return 0xD3;
    case VK_PRIOR: return 0xC9; case VK_NEXT: return 0xD1; case VK_PAUSE: return 0xC5; case VK_SCROLL: return 0x46;
    case VK_F1: return 0x3B; case VK_F2: return 0x3C; case VK_F3: return 0x3D; case VK_F4: return 0x3E; case VK_F5: return 0x3F;
    case VK_F6: return 0x40; case VK_F7: return 0x41; case VK_F8: return 0x42; case VK_F9: return 0x43; case VK_F10: return 0x44;
    case VK_F11: return 0x57; case VK_F12: return 0x58; case VK_OEM_3: return 0x29; case VK_TAB: return 0x0F;
    default: return 0;
    }
}

// Polled input. We look at the real state first (for the toggle key), then hand the game a buffer
// full of zeros while the overlay owns the input. cb == 256 is the keyboard's state array.
HRESULT STDMETHODCALLTYPE GetDeviceStateHook(void* self, DWORD cb, LPVOID data)
{
    HRESULT hr = g_origGetDeviceState(self, cb, data);
    if (SUCCEEDED(hr) && data && cb == 256 && IsBlockedDevice(self)) {
        int dik = VkToDik(g_cfg.overlayKey);
        if (dik) { g_dikSeen = true; g_dikKeyDown = (static_cast<const BYTE*>(data)[dik] & 0x80) != 0; }
    }
    if (OverlayOwnsInput() && IsBlockedDevice(self) && data && cb) { memset(data, 0, cb); return DI_OK; }
    return hr;
}

// Buffered input, same idea. The original is still called while the overlay is open - that drains
// the buffer, so the game does not get a burst of stale key presses when the overlay closes.
HRESULT STDMETHODCALLTYPE GetDeviceDataHook(void* self, DWORD cbObj, void* rgdod, LPDWORD pdwInOut, DWORD flags)
{
    HRESULT hr = g_origGetDeviceData(self, cbObj, rgdod, pdwInOut, flags);
    if (SUCCEEDED(hr) && rgdod && pdwInOut && IsBlockedDevice(self) && cbObj >= 8) {
        int dik = VkToDik(g_cfg.overlayKey);
        const BYTE* e = static_cast<const BYTE*>(rgdod);
        for (DWORD i = 0; dik && i < *pdwInOut; ++i, e += cbObj) {
            DWORD ofs = *reinterpret_cast<const DWORD*>(e), data = *reinterpret_cast<const DWORD*>(e + 4);
            if (static_cast<int>(ofs) == dik) { g_dikSeen = true; g_dikKeyDown = (data & 0x80) != 0; }
        }
    }
    if (OverlayOwnsInput() && IsBlockedDevice(self)) { if (pdwInOut) *pdwInOut = 0; return DI_OK; }
    return hr;
}

// Remembers the keyboard and mouse devices the game creates. Gamepads and wheels are left alone.
HRESULT STDMETHODCALLTYPE DICreateDeviceHook(void* self, REFGUID rguid, void** out, LPUNKNOWN outer)
{
    HRESULT hr = g_origDICreateDevice(self, rguid, out, outer);
    if (FAILED(hr) || !out || !*out) return hr;
    static const GUID kSysKeyboard = { 0x6F1D2B61, 0xD5A0, 0x11CF, { 0xBF, 0xC7, 0x44, 0x45, 0x53, 0x54, 0x00, 0x00 } };
    static const GUID kSysMouse    = { 0x6F1D2B60, 0xD5A0, 0x11CF, { 0xBF, 0xC7, 0x44, 0x45, 0x53, 0x54, 0x00, 0x00 } };
    if (!IsEqualGUID(rguid, kSysKeyboard) && !IsEqualGUID(rguid, kSysMouse)) return hr;
    if (g_blockedCount < 8) g_blockedDevices[g_blockedCount++] = *out;
    void** vt = *reinterpret_cast<void***>(*out);   // slot 9 = GetDeviceState, 10 = GetDeviceData
    if (vt[9] != reinterpret_cast<void*>(&GetDeviceStateHook)) {
        g_origGetDeviceState = reinterpret_cast<DIGetDeviceState_t>(SwapVtableEntry(*out, 9,  reinterpret_cast<void*>(&GetDeviceStateHook)));
        g_origGetDeviceData  = reinterpret_cast<DIGetDeviceData_t>(SwapVtableEntry(*out, 10, reinterpret_cast<void*>(&GetDeviceDataHook)));
    }
    PluginLog("input: keyboard/mouse device created, blocked while the overlay is open");
    return hr;
}

HRESULT WINAPI DirectInput8CreateHook(HINSTANCE h, DWORD ver, REFIID riid, LPVOID* out, LPUNKNOWN outer)
{
    HRESULT hr = g_origDI8Create(h, ver, riid, out, outer);
    if (SUCCEEDED(hr) && out && *out && !g_origDICreateDevice)
        g_origDICreateDevice = reinterpret_cast<DICreateDevice_t>(SwapVtableEntry(*out, 3, reinterpret_cast<void*>(&DICreateDeviceHook)));
    return hr;
}

// --- Window / ImGui ---------------------------------------------------
// While the overlay is open, ImGui gets the window messages and the game does not see mouse or keys.
LRESULT CALLBACK WndProcHook(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    if (OverlayOwnsInput()) {
        ImGui_ImplWin32_WndProcHandler(hWnd, msg, wParam, lParam);
        switch (msg) {
        case WM_MOUSEMOVE: case WM_LBUTTONDOWN: case WM_LBUTTONUP: case WM_LBUTTONDBLCLK:
        case WM_RBUTTONDOWN: case WM_RBUTTONUP: case WM_MBUTTONDOWN: case WM_MBUTTONUP:
        case WM_MOUSEWHEEL: case WM_KEYDOWN: case WM_KEYUP: case WM_SYSKEYDOWN: case WM_SYSKEYUP: case WM_CHAR:
            return 0;
        }
    }
    return CallWindowProc(g_origWndProc, hWnd, msg, wParam, lParam);
}

// Set up lazily, the first time the overlay is opened - whoever never presses the key never pays for it.
void InitImGui(IDirect3DDevice9* device)
{
    D3DDEVICE_CREATION_PARAMETERS cp = {};
    device->GetCreationParameters(&cp);
    g_hWnd = cp.hFocusWindow;
    if (!g_hWnd) return;
    ImGui::CreateContext();
    ImGuiIO& io = ImGui::GetIO();
    io.IniFilename = nullptr;       // no imgui.ini in the game folder
    io.MouseDrawCursor = true;      // the game hides the system cursor
    ImGui::StyleColorsDark();
    ImGui::GetStyle().WindowRounding = 4.f;
    ImGui_ImplWin32_Init(g_hWnd);
    ImGui_ImplDX9_Init(device);
    g_origWndProc = reinterpret_cast<WNDPROC>(SetWindowLongPtr(g_hWnd, GWLP_WNDPROC, reinterpret_cast<LONG_PTR>(WndProcHook)));
    g_imguiReady = true;
    PluginLog("overlay: ImGui initialised (hwnd %p)", g_hWnd);
}

#include "overlay_ui.inc"

// --- D3D9 ---------------------------------------------------------------
// Usage tracking costs frame time, so it only runs while the overlay is open. And closing the
// overlay must not leave a hover checkerboard behind.
void SetVisible(bool visible)
{
    g_visible = visible;
    TexturesSetTracking(visible);
    if (!visible) TexturesSetHover(0);
}

HRESULT APIENTRY EndSceneHook(IDirect3DDevice9* device)
{
    // Prefer what the game's own DirectInput keyboard reported: GetAsyncKeyState can miss keys while
    // the game holds the keyboard exclusively.
    bool down = g_cfg.overlay && (g_dikSeen ? g_dikKeyDown : (GetAsyncKeyState(g_cfg.overlayKey) & 0x8000) != 0);
    if (down && !g_keyWasDown) SetVisible(!g_visible);
    g_keyWasDown = down;
    TexturesEndFrame();
    if (!g_visible) return g_origEndScene(device);      // closed: nothing else runs

    ++g_frames;
    DWORD now = GetTickCount();
    if (now - g_lastTick >= 500) { g_fps = g_frames * 1000.f / (now - g_lastTick); g_frames = 0; g_lastTick = now; }
    if (!g_imguiReady) InitImGui(device);
    if (g_imguiReady) {
        ImGui_ImplDX9_NewFrame();
        ImGui_ImplWin32_NewFrame();
        ImGui::NewFrame();
        DrawUI();
        if (!g_visible) SetVisible(false);              // closed with the window's X rather than the key
        ImGui::EndFrame();
        ImGui::Render();
        // The thumbnails must show the textures themselves, not what SetTexture would swap them for -
        // otherwise "game" and "replacement" in the detail panel would be the same picture.
        TexturesBypassBegin();
        ImGui_ImplDX9_RenderDrawData(ImGui::GetDrawData());
        TexturesBypassEnd();
    }
    return g_origEndScene(device);
}

// ImGui's font texture and buffers live in the default pool and have to be let go around a Reset
// (alt-tab, resolution change), or the Reset fails.
HRESULT APIENTRY ResetHook(IDirect3DDevice9* device, D3DPRESENT_PARAMETERS* pp)
{
    if (g_imguiReady) ImGui_ImplDX9_InvalidateDeviceObjects();
    HRESULT hr = g_origReset(device, pp);
    if (g_imguiReady && SUCCEEDED(hr)) ImGui_ImplDX9_CreateDeviceObjects();
    return hr;
}

// Device vtable slots: 16 = Reset, 42 = EndScene.
void HookDeviceVTable(IDirect3DDevice9* device)
{
    if (g_deviceHooked || !device) return;
    void** vt = *reinterpret_cast<void***>(device);
    if (MH_CreateHook(vt[42], &EndSceneHook, reinterpret_cast<void**>(&g_origEndScene)) != MH_OK ||
        MH_CreateHook(vt[16], &ResetHook,    reinterpret_cast<void**>(&g_origReset))    != MH_OK ||
        MH_EnableHook(vt[42]) != MH_OK || MH_EnableHook(vt[16]) != MH_OK) {
        PluginLog("overlay: hooking EndScene/Reset failed");
        return;
    }
    g_deviceHooked = true;
    TexturesOnDevice(device);
    PluginLog("overlay: EndScene/Reset hooked at device creation (device %p), toggle key 0x%02X", device, g_cfg.overlayKey);
}

HRESULT APIENTRY CreateDeviceHook(IDirect3D9* self, UINT adapter, D3DDEVTYPE type, HWND focus, DWORD behaviour,
                                  D3DPRESENT_PARAMETERS* pp, IDirect3DDevice9** out)
{
    HRESULT hr = g_origCreateDevice(self, adapter, type, focus, behaviour, pp, out);
    if (SUCCEEDED(hr) && out && *out) HookDeviceVTable(*out);
    return hr;
}

IDirect3D9* WINAPI Direct3DCreate9Hook(UINT sdkVersion)
{
    IDirect3D9* d3d = g_origDirect3DCreate9(sdkVersion);
    if (d3d && !g_origCreateDevice) {   // IDirect3D9 slot 16 = CreateDevice
        g_origCreateDevice = reinterpret_cast<CreateDevice_t>(SwapVtableEntry(d3d, 16, reinterpret_cast<void*>(&CreateDeviceHook)));
        PluginLog("overlay: IDirect3D9::CreateDevice captured");
    }
    return d3d;
}

} // namespace

void OverlayInit()
{
    MH_STATUS st = MH_Initialize();
    if (st != MH_OK && st != MH_ERROR_ALREADY_INITIALIZED) { PluginLog("overlay: MH_Initialize failed"); return; }
    g_origDirect3DCreate9 = reinterpret_cast<Direct3DCreate9_t>(PatchImport("d3d9.dll", "Direct3DCreate9", reinterpret_cast<void*>(&Direct3DCreate9Hook)));
    PluginLog(g_origDirect3DCreate9 ? "overlay: Direct3DCreate9 import redirected" : "overlay: Direct3DCreate9 import not found, overlay disabled");
    // With the overlay switched off in the INI the device hooks above are still needed: they drive
    // auto apply and hot reload. Only the key and the input blocking go away.
    if (g_cfg.overlay && g_cfg.inputBlock) {
        g_origDI8Create = reinterpret_cast<DI8Create_t>(PatchImport("dinput8.dll", "DirectInput8Create", reinterpret_cast<void*>(&DirectInput8CreateHook)));
        PluginLog(g_origDI8Create ? "input: DirectInput8Create import redirected" : "input: DirectInput8Create import not found");
    }
}