2077-gps-mod/red4ext/EdgeWeightGPS/src/Main.cpp

#include <RED4ext/Api/ApiVersion.hpp>
#include <RED4ext/Api/v1/EMainReason.hpp>
#include <RED4ext/Api/v1/PluginHandle.hpp>
#include <RED4ext/Api/v1/PluginInfo.hpp>
#include <RED4ext/Api/v1/Runtime.hpp>
#include <RED4ext/Api/v1/Version.hpp>
#include <RED4ext/Common.hpp>

#include <windows.h>

#include <array>
#include <cstdint>
#include <cstring>
#include <initializer_list>
#include <sstream>
#include <string>

namespace RED4ext::v1
{
struct Logger
{
    void (*Trace)(PluginHandle aHandle, const char* aMessage);
    void (*TraceF)(PluginHandle aHandle, const char* aFormat, ...);
    void (*TraceW)(PluginHandle aHandle, const wchar_t* aMessage);
    void (*TraceWF)(PluginHandle aHandle, const wchar_t* aFormat, ...);
    void (*Debug)(PluginHandle aHandle, const char* aMessage);
    void (*DebugF)(PluginHandle aHandle, const char* aFormat, ...);
    void (*DebugW)(PluginHandle aHandle, const wchar_t* aMessage);
    void (*DebugWF)(PluginHandle aHandle, const wchar_t* aFormat, ...);
    void (*Info)(PluginHandle aHandle, const char* aMessage);
    void (*InfoF)(PluginHandle aHandle, const char* aFormat, ...);
    void (*InfoW)(PluginHandle aHandle, const wchar_t* aMessage);
    void (*InfoWF)(PluginHandle aHandle, const wchar_t* aFormat, ...);
    void (*Warn)(PluginHandle aHandle, const char* aMessage);
    void (*WarnF)(PluginHandle aHandle, const char* aFormat, ...);
    void (*WarnW)(PluginHandle aHandle, const wchar_t* aMessage);
    void (*WarnWF)(PluginHandle aHandle, const wchar_t* aFormat, ...);
    void (*Error)(PluginHandle aHandle, const char* aMessage);
    void (*ErrorF)(PluginHandle aHandle, const char* aFormat, ...);
    void (*ErrorW)(PluginHandle aHandle, const wchar_t* aMessage);
    void (*ErrorWF)(PluginHandle aHandle, const wchar_t* aFormat, ...);
    void (*Critical)(PluginHandle aHandle, const char* aMessage);
    void (*CriticalF)(PluginHandle aHandle, const char* aFormat, ...);
    void (*CriticalW)(PluginHandle aHandle, const wchar_t* aMessage);
    void (*CriticalWF)(PluginHandle aHandle, const wchar_t* aFormat, ...);
};

struct Hooking
{
    bool (*Attach)(PluginHandle aHandle, void* aTarget, void* aDetour, void** aOriginal);
    bool (*Detach)(PluginHandle aHandle, void* aTarget);
};

struct GameState
{
    bool (*OnEnter)(void* aApp);
    bool (*OnUpdate)(void* aApp);
    bool (*OnExit)(void* aApp);
};

struct GameStates
{
    bool (*Add)(PluginHandle aHandle, uint32_t aType, GameState* aState);
};

struct Scripts
{
    bool (*Add)(PluginHandle aHandle, const wchar_t* aPath);
    bool (*RegisterNeverRefType)(const char* aType);
    bool (*RegisterMixedRefType)(const char* aType);
};

struct Sdk
{
    SemVer* runtime;
    Logger* logger;
    Hooking* hooking;
    GameStates* gameStates;
    Scripts* scripts;
};
}

namespace
{
constexpr uintptr_t kGpsRelaxUpdateRva = 0x040BA58;
constexpr uintptr_t kGpsExpandListARelaxPatchRva = 0x040B49A;
constexpr uintptr_t kGpsExpandListBRelaxPatchRva = 0x040B6C3;
constexpr size_t kRelaxPatchSize = 23;
constexpr size_t kPatchCaveSize = 80;
constexpr float kFixedEdgePenalty = 80.0f;

struct PatchSite
{
    const char* name;
    uintptr_t targetRva;
    std::array<uint8_t, 5> expectedFirstInstruction;
    std::array<uint8_t, kRelaxPatchSize> original{};
    uint8_t* cave = nullptr;
    bool applied = false;
};

HMODULE gModule = nullptr;
RED4ext::v1::PluginHandle gHandle = nullptr;
const RED4ext::v1::Sdk* gSdk = nullptr;

std::array<PatchSite, 2> gPatchSites{{
    {
        "expand-list-a",
        kGpsExpandListARelaxPatchRva,
        {0xF3, 0x41, 0x0F, 0x58, 0xF6}, // addss xmm6, xmm14
    },
    {
        "expand-list-b",
        kGpsExpandListBRelaxPatchRva,
        {0xF3, 0x41, 0x0F, 0x58, 0xF4}, // addss xmm6, xmm12
    },
}};

uintptr_t GameImageBase()
{
    return reinterpret_cast<uintptr_t>(GetModuleHandleW(nullptr));
}

void LogInfo(const std::string& aMessage)
{
    if (gSdk && gSdk->logger && gSdk->logger->Info)
    {
        gSdk->logger->Info(gHandle, aMessage.c_str());
    }
}

void LogError(const std::string& aMessage)
{
    if (gSdk && gSdk->logger && gSdk->logger->Error)
    {
        gSdk->logger->Error(gHandle, aMessage.c_str());
        return;
    }

    LogInfo(aMessage);
}

const char* ReasonToString(RED4ext::v1::EMainReason aReason)
{
    switch (aReason)
    {
    case RED4ext::v1::EMainReason::Load:
        return "Load";
    case RED4ext::v1::EMainReason::Unload:
        return "Unload";
    default:
        return "Unknown";
    }
}

void WriteInt32(uint8_t* aAddress, int32_t aValue)
{
    std::memcpy(aAddress, &aValue, sizeof(aValue));
}

void WriteFloat(uint8_t* aAddress, float aValue)
{
    std::memcpy(aAddress, &aValue, sizeof(aValue));
}

void EmitBytes(uint8_t* aCave, size_t& aOffset, std::initializer_list<uint8_t> aBytes)
{
    for (auto byte : aBytes)
    {
        aCave[aOffset++] = byte;
    }
}

void EmitAbsoluteRaxCall(uint8_t* aCave, size_t& aOffset, uintptr_t aTarget)
{
    EmitBytes(aCave, aOffset, {0x48, 0xB8}); // mov rax, imm64
    std::memcpy(aCave + aOffset, &aTarget, sizeof(aTarget));
    aOffset += sizeof(aTarget);
    EmitBytes(aCave, aOffset, {0xFF, 0xD0}); // call rax
}

void EmitAbsoluteRaxJump(uint8_t* aCave, size_t& aOffset, uintptr_t aTarget)
{
    EmitBytes(aCave, aOffset, {0x48, 0xB8}); // mov rax, imm64
    std::memcpy(aCave + aOffset, &aTarget, sizeof(aTarget));
    aOffset += sizeof(aTarget);
    EmitBytes(aCave, aOffset, {0xFF, 0xE0}); // jmp rax
}

void BuildAbsoluteJumpPatch(uint8_t* aPatch, size_t aPatchSize, uintptr_t aTarget)
{
    size_t offset = 0;
    aPatch[offset++] = 0xFF; // jmp qword ptr [rip+0]
    aPatch[offset++] = 0x25;
    aPatch[offset++] = 0x00;
    aPatch[offset++] = 0x00;
    aPatch[offset++] = 0x00;
    aPatch[offset++] = 0x00;
    std::memcpy(aPatch + offset, &aTarget, sizeof(aTarget));
    offset += sizeof(aTarget);

    while (offset < aPatchSize)
    {
        aPatch[offset++] = 0x90;
    }
}

bool InstallPatchSite(PatchSite& aSite)
{
    if (aSite.applied)
    {
        return true;
    }

    auto* target = reinterpret_cast<uint8_t*>(GameImageBase() + aSite.targetRva);
    if (std::memcmp(target, aSite.expectedFirstInstruction.data(), aSite.expectedFirstInstruction.size()) != 0)
    {
        std::ostringstream line;
        line << "GPS fixed-edge penalty skipped site=" << aSite.name
             << " reason=unexpected-bytes rva=0x" << std::hex << aSite.targetRva;
        LogError(line.str());
        return false;
    }

    auto* cave = static_cast<uint8_t*>(VirtualAlloc(nullptr, kPatchCaveSize, MEM_COMMIT | MEM_RESERVE,
                                                   PAGE_EXECUTE_READWRITE));
    if (!cave)
    {
        std::ostringstream line;
        line << "GPS fixed-edge penalty skipped site=" << aSite.name << " reason=alloc-failed";
        LogError(line.str());
        return false;
    }

    size_t offset = 0;
    for (auto byte : aSite.expectedFirstInstruction)
    {
        cave[offset++] = byte;
    }

    const auto addPenaltyOffset = offset;
    EmitBytes(cave, offset, {0xF3, 0x0F, 0x58, 0x35, 0x00, 0x00, 0x00, 0x00}); // addss xmm6, [rip+disp32]
    EmitBytes(cave, offset, {0xF3, 0x0F, 0x58, 0xC6});                         // addss xmm0, xmm6
    EmitBytes(cave, offset, {0x0F, 0x28, 0xDE});                               // movaps xmm3, xmm6
    EmitBytes(cave, offset, {0xF3, 0x0F, 0x11, 0x44, 0x24, 0x20});             // movss [rsp+0x20], xmm0
    EmitAbsoluteRaxCall(cave, offset, GameImageBase() + kGpsRelaxUpdateRva);
    EmitAbsoluteRaxJump(cave, offset, GameImageBase() + aSite.targetRva + kRelaxPatchSize);

    while ((offset % alignof(float)) != 0)
    {
        cave[offset++] = 0x90;
    }

    const auto penaltyConstantOffset = offset;
    WriteFloat(cave + offset, kFixedEdgePenalty);
    offset += sizeof(float);

    const auto caveAddress = reinterpret_cast<uintptr_t>(cave);
    WriteInt32(cave + addPenaltyOffset + 4,
               static_cast<int32_t>((caveAddress + penaltyConstantOffset) - (caveAddress + addPenaltyOffset + 8)));

    std::memcpy(aSite.original.data(), target, aSite.original.size());

    std::array<uint8_t, kRelaxPatchSize> patch{};
    BuildAbsoluteJumpPatch(patch.data(), patch.size(), caveAddress);

    DWORD oldProtect = 0;
    if (!VirtualProtect(target, patch.size(), PAGE_EXECUTE_READWRITE, &oldProtect))
    {
        VirtualFree(cave, 0, MEM_RELEASE);
        std::ostringstream line;
        line << "GPS fixed-edge penalty skipped site=" << aSite.name << " reason=protect-failed";
        LogError(line.str());
        return false;
    }

    std::memcpy(target, patch.data(), patch.size());
    FlushInstructionCache(GetCurrentProcess(), target, patch.size());

    DWORD ignoredProtect = 0;
    VirtualProtect(target, patch.size(), oldProtect, &ignoredProtect);

    aSite.cave = cave;
    aSite.applied = true;

    std::ostringstream line;
    line << "GPS fixed-edge penalty applied site=" << aSite.name << " rva=0x" << std::hex << aSite.targetRva
         << " cave=" << static_cast<void*>(cave) << std::dec << " penalty=" << kFixedEdgePenalty;
    LogInfo(line.str());
    return true;
}

void RemovePatchSites()
{
    for (auto& site : gPatchSites)
    {
        if (!site.applied)
        {
            continue;
        }

        auto* target = reinterpret_cast<uint8_t*>(GameImageBase() + site.targetRva);
        DWORD oldProtect = 0;
        if (VirtualProtect(target, site.original.size(), PAGE_EXECUTE_READWRITE, &oldProtect))
        {
            std::memcpy(target, site.original.data(), site.original.size());
            FlushInstructionCache(GetCurrentProcess(), target, site.original.size());

            DWORD ignoredProtect = 0;
            VirtualProtect(target, site.original.size(), oldProtect, &ignoredProtect);

            std::ostringstream line;
            line << "GPS fixed-edge penalty restored site=" << site.name;
            LogInfo(line.str());
        }
        else
        {
            std::ostringstream line;
            line << "GPS fixed-edge penalty restore failed site=" << site.name << " reason=protect-failed";
            LogError(line.str());
        }

        if (site.cave)
        {
            VirtualFree(site.cave, 0, MEM_RELEASE);
            site.cave = nullptr;
        }

        site.applied = false;
    }
}

bool InstallPatchSites()
{
    bool ok = true;
    for (auto& site : gPatchSites)
    {
        ok = InstallPatchSite(site) && ok;
    }

    if (!ok)
    {
        RemovePatchSites();
    }

    return ok;
}
}

BOOL APIENTRY DllMain(HMODULE aModule, DWORD aReason, LPVOID)
{
    if (aReason == DLL_PROCESS_ATTACH)
    {
        gModule = aModule;
        DisableThreadLibraryCalls(aModule);
    }

    return TRUE;
}

RED4EXT_C_EXPORT bool RED4EXT_CALL Main(RED4ext::v1::PluginHandle aHandle, RED4ext::v1::EMainReason aReason,
                                        const RED4ext::v1::Sdk* aSdk)
{
    gHandle = aHandle;
    gSdk = aSdk;

    std::ostringstream line;
    line << "MomentumGPS Main reason=" << ReasonToString(aReason) << " module=" << gModule;
    if (aSdk && aSdk->runtime)
    {
        line << " runtime=" << static_cast<uint32_t>(aSdk->runtime->major) << "." << aSdk->runtime->minor << "."
             << aSdk->runtime->patch;
    }
    LogInfo(line.str());

    if (aReason == RED4ext::v1::EMainReason::Load)
    {
        InstallPatchSites();
    }
    else if (aReason == RED4ext::v1::EMainReason::Unload)
    {
        RemovePatchSites();
    }

    return true;
}

RED4EXT_C_EXPORT void RED4EXT_CALL Query(RED4ext::v1::PluginInfo* aInfo)
{
    aInfo->name = L"MomentumGPS";
    aInfo->author = L"salt+Codex";
    aInfo->version = RED4ext::v1::SemVer{0, 2, 0, {0, 0}};
    aInfo->runtime = RED4ext::v1::FileVer{3, 0, 80, 51928};
    aInfo->sdk = RED4ext::v1::SemVer{0, 5, 0, {0, 0}};
}

RED4EXT_C_EXPORT uint32_t RED4EXT_CALL Supports()
{
    return RED4EXT_API_VERSION_1_COMPAT_0;
}