#include "antihookingprotection.h"

#include <NktHookLib.h>

typedef HMODULE (WINAPI *LoadLibraryAFunc)(LPCSTR lpLibFileName);
typedef HMODULE (WINAPI *LoadLibraryWFunc)(LPCWSTR lpLibFileName);
typedef HMODULE (WINAPI *LoadLibraryExAFunc)(LPCSTR lpLibFileName, HANDLE hFile, DWORD dwFlags);
typedef HMODULE (WINAPI *LoadLibraryExWFunc)(LPCWSTR lpLibFileName, HANDLE hFile, DWORD dwFlags);

class AntiHookingProtection
{
public:
    static void enable()
    {
    #ifdef QT_DEBUG
        s_HookManager.SetEnableDebugOutput(true);
    #endif

        HINSTANCE kernel32Handle = NktHookLibHelpers::GetModuleBaseAddress(L"kernel32.dll");
        SIZE_T hookId;

        s_HookManager.Hook(&hookId, (LPVOID*)&s_RealLoadLibraryA,
                           NktHookLibHelpers::GetProcedureAddress(kernel32Handle, "LoadLibraryA"),
                           (LPVOID)AntiHookingProtection::LoadLibraryAHook);
        s_HookManager.Hook(&hookId, (LPVOID*)&s_RealLoadLibraryW,
                           NktHookLibHelpers::GetProcedureAddress(kernel32Handle, "LoadLibraryW"),
                           (LPVOID)AntiHookingProtection::LoadLibraryWHook);
        s_HookManager.Hook(&hookId, (LPVOID*)&s_RealLoadLibraryExA,
                           NktHookLibHelpers::GetProcedureAddress(kernel32Handle, "LoadLibraryExA"),
                           (LPVOID)AntiHookingProtection::LoadLibraryExAHook);
        s_HookManager.Hook(&hookId, (LPVOID*)&s_RealLoadLibraryExW,
                           NktHookLibHelpers::GetProcedureAddress(kernel32Handle, "LoadLibraryExW"),
                           (LPVOID)AntiHookingProtection::LoadLibraryExWHook);
    }

private:
    static bool isImageBlacklistedW(LPCWSTR lpLibFileName)
    {
        LPCWSTR dllName;

        // If the library has a path prefixed, remove it
        dllName = wcsrchr(lpLibFileName, '\\');
        if (!dllName) {
            // No prefix, so use the full name
            dllName = lpLibFileName;
        }
        else {
            // Advance past the backslash
            dllName++;
        }

        // FIXME: We don't currently handle LoadLibrary calls where the
        // library name does not include a file extension and the loader
        // automatically assumes .dll.

        for (int i = 0; i < ARRAYSIZE(k_BlacklistedDlls); i++) {
            if (_wcsicmp(dllName, k_BlacklistedDlls[i]) == 0) {
                return true;
            }
        }

        return false;
    }

    static bool isImageBlacklistedA(LPCSTR lpLibFileName)
    {
        int uniChars = MultiByteToWideChar(CP_THREAD_ACP, 0, lpLibFileName, -1, nullptr, 0);
        if (uniChars > 0) {
            PWCHAR wideBuffer = new WCHAR[uniChars];
            uniChars = MultiByteToWideChar(CP_THREAD_ACP, 0,
                                           lpLibFileName, -1,
                                           wideBuffer, uniChars * sizeof(WCHAR));
            if (uniChars > 0) {
                bool ret = isImageBlacklistedW(wideBuffer);
                delete[] wideBuffer;
                return ret;
            }
            else {
                delete[] wideBuffer;
            }
        }

        // Error path
        return false;
    }

    static HMODULE WINAPI LoadLibraryAHook(LPCSTR lpLibFileName)
    {
        if (lpLibFileName && isImageBlacklistedA(lpLibFileName)) {
            SetLastError(ERROR_ACCESS_DISABLED_BY_POLICY);
            return nullptr;
        }

        return s_RealLoadLibraryA(lpLibFileName);
    }

    static HMODULE WINAPI LoadLibraryWHook(LPCWSTR lpLibFileName)
    {
        if (lpLibFileName && isImageBlacklistedW(lpLibFileName)) {
            SetLastError(ERROR_ACCESS_DISABLED_BY_POLICY);
            return nullptr;
        }

        return s_RealLoadLibraryW(lpLibFileName);
    }

    static HMODULE WINAPI LoadLibraryExAHook(LPCSTR lpLibFileName, HANDLE hFile, DWORD dwFlags)
    {
        if (lpLibFileName && isImageBlacklistedA(lpLibFileName)) {
            SetLastError(ERROR_ACCESS_DISABLED_BY_POLICY);
            return nullptr;
        }

        return s_RealLoadLibraryExA(lpLibFileName, hFile, dwFlags);
    }

    static HMODULE WINAPI LoadLibraryExWHook(LPCWSTR lpLibFileName, HANDLE hFile, DWORD dwFlags)
    {
        if (lpLibFileName && isImageBlacklistedW(lpLibFileName)) {
            SetLastError(ERROR_ACCESS_DISABLED_BY_POLICY);
            return nullptr;
        }

        return s_RealLoadLibraryExW(lpLibFileName, hFile, dwFlags);
    }

    static CNktHookLib s_HookManager;
    static LoadLibraryAFunc s_RealLoadLibraryA;
    static LoadLibraryWFunc s_RealLoadLibraryW;
    static LoadLibraryExAFunc s_RealLoadLibraryExA;
    static LoadLibraryExWFunc s_RealLoadLibraryExW;

    static constexpr LPCWSTR k_BlacklistedDlls[] = {
        // These A-Volute DLLs shipped with various audio driver packages improperly handle
        // D3D9 exclusive fullscreen in a way that causes CreateDeviceEx() to deadlock.
        // https://github.com/moonlight-stream/moonlight-qt/issues/102
        L"NahimicOSD.dll", // ASUS Sonic Radar 3
        L"SSAudioOSD.dll", // SteelSeries headsets
        L"SS2OSD.dll", // ASUS Sonic Studio 2
        L"Nahimic2OSD.dll",
        L"NahimicMSIOSD.dll",
        L"nhAsusPhoebusOSD.dll" // ASUS Phoebus

        // This DLL has been seen in several crash reports. Some Googling
        // suggests it's highly unstable and causes issues in many games.
        L"EZFRD32.dll",
        L"EZFRD64.dll",

        // These are the newer dList DLLs for Optimus hybrid graphics DDI.
        // https://docs.microsoft.com/en-us/windows-hardware/drivers/display/hybrid-system-ddi
        //
        // We forcefully block them from loading because Optimus has a bug that
        // deadlocks DXVA2 when we present with D3DPRESENT_DONOTWAIT. This will prevent
        // Optimus from ever using the dGPU even if the user has requested it.
        // https://github.com/moonlight-stream/moonlight-qt/issues/240
        // https://github.com/moonlight-stream/moonlight-qt/issues/235
        L"nvdlist.dll",
        L"nvdlistx.dll",

        // These are the older dList/AppInit DLLs for Optimus hybrid graphics DDI.
        // https://docs.microsoft.com/en-us/windows-hardware/drivers/display/hybrid-system-ddi
        //
        // These seem to cause a crash in PresentEx() in full-screen exclusive mode.
        // This block will prevent Optimus from ever using the dGPU even if the user has requested it.
        // https://github.com/moonlight-stream/moonlight-qt/issues/386
        //
        // d3d9!CSwapChain::BltToHybridPrimary+0x200:
        // 00007ffa`23f37e58 488b01          mov     rax,qword ptr [rcx] ds:00000000`00000038=????????????????
        // 00 0000004e`496ff4e0 00007ffa`23f39e2c d3d9!CSwapChain::BltToHybridPrimary+0x200
        // 01 0000004e`496ff880 00007ffa`23ee39ce d3d9!CSwapChain::FlipToSurface+0x15c
        // 02 0000004e`496ff900 00007ffa`23f4dd75 d3d9!CSwapChain::PresentMain+0x3e13e
        // 03 0000004e`496ffab0 00007ffa`23f4dccd d3d9!CBaseDevice::PresentMain+0x9d
        // 04 0000004e`496ffb00 00007ff7`8e31016f d3d9!CBaseDevice::PresentEx+0xbd
        // 05 0000004e`496ffb50 00007ff7`8e30df1e Moonlight!DXVA2Renderer::renderFrame+0x61f [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\dxva2.cpp @ 1035]
        // 06 0000004e`496ffd50 00007ff7`8e30e46a Moonlight!Pacer::renderFrame+0x3e [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\pacer\pacer.cpp @ 265]
        // 07 (Inline Function) --------`-------- Moonlight!Pacer::renderLastFrameAndUnlock+0x197 [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\pacer\pacer.cpp @ 156]
        // 08 0000004e`496ffda0 00007ffa`14476978 Moonlight!Pacer::renderThread+0x1ca [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\pacer\pacer.cpp @ 88]
        // 09 0000004e`496ffde0 00007ffa`14476ee2 SDL2!SDL_RunThread+0x38 [C:\Users\camer\SDL\src\thread\SDL_thread.c @ 276]
        // 0a 0000004e`496ffe10 00007ffa`2aae0e82 SDL2!RunThread+0x12 [C:\Users\camer\SDL\src\thread\windows\SDL_systhread.c @ 83]
        // 0b 0000004e`496ffe40 00007ffa`2d627bd4 ucrtbase!thread_start<unsigned int (__cdecl*)(void *),1>+0x42
        // 0c 0000004e`496ffe70 00007ffa`2da4ce51 kernel32!BaseThreadInitThunk+0x14
        // 0d 0000004e`496ffea0 00000000`00000000 ntdll!RtlUserThreadStart+0x21
        //
        L"nvinit.dll",
        L"nvinitx.dll",

        // In some unknown circumstances, RTSS tries to hook in the middle of an instruction, leaving garbage
        // code inside d3d9.dll that causes a crash when executed:
        //
        // 0:000> u
        // d3d9!D3D9GetCurrentOwnershipMode+0x5d:
        // 00007ff8`95b95861 9b              wait
        // 00007ff8`95b95862 a7              cmps    dword ptr [rsi],dword ptr [rdi]   <--- crash happens here
        // 00007ff8`95b95863 ff              ???
        // 00007ff8`95b95864 bfe8ca8a00      mov     edi,8ACAE8h
        // 00007ff8`95b95869 00eb            add     bl,ch
        // 00007ff8`95b9586b f1              ???
        // 00007ff8`95b9586c b808000000      mov     eax,8
        // 00007ff8`95b95871 ebe6            jmp     d3d9!D3D9GetCurrentOwnershipMode+0x55 (00007ff8`95b95859)
        //
        // Disassembling starting at the exact address of the attempted hook yields the intended jmp instruction
        //
        // 0:000> u d3d9!D3D9GetCurrentOwnershipMode+0x5c:
        // 00007ff8`95b95860 e99ba7ffbf      jmp     00007ff8`55b90000
        //
        // Since the RTSS OSD doesn't even work with DXVA2, we'll just block the hooks entirely.
        L"RTSSHooks.dll",
        L"RTSSHooks64.dll",
    };
};

CNktHookLib AntiHookingProtection::s_HookManager;
LoadLibraryAFunc AntiHookingProtection::s_RealLoadLibraryA;
LoadLibraryWFunc AntiHookingProtection::s_RealLoadLibraryW;
LoadLibraryExAFunc AntiHookingProtection::s_RealLoadLibraryExA;
LoadLibraryExWFunc AntiHookingProtection::s_RealLoadLibraryExW;

AH_EXPORT void AntiHookingDummyImport() {}

extern "C"
BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID)
{
    switch (fdwReason)
    {
    case DLL_PROCESS_ATTACH:
        AntiHookingProtection::enable();
        DisableThreadLibraryCalls(hinstDLL);
        break;
    }

    return TRUE;
};