diff --git a/include/RED4ext/Relocation.hpp b/include/RED4ext/Relocation.hpp index 487c11efa..8f8469654 100644 --- a/include/RED4ext/Relocation.hpp +++ b/include/RED4ext/Relocation.hpp @@ -88,12 +88,19 @@ class RelocVtbl : private RelocBase uintptr_t* m_address; }; +enum class UniversalRelocSegment +{ + Text = 0, + Data, + Rdata +}; + class UniversalRelocBase { public: - uintptr_t Resolve(uint64_t aHash) + static uintptr_t Resolve(UniversalRelocSegment aSegment, uint64_t aHash) { - using functionType = uintptr_t (*)(uint64_t); + using functionType = uintptr_t (*)(UniversalRelocSegment, uint64_t); static functionType resolveFunc = nullptr; static std::once_flag flag; @@ -114,8 +121,7 @@ class UniversalRelocBase TerminateProcess(GetCurrentProcess(), 1); } - resolveFunc = - reinterpret_cast(GetProcAddress(handle, functionName)); + resolveFunc = reinterpret_cast(GetProcAddress(handle, functionName)); if (resolveFunc == nullptr) { std::stringstream stream; @@ -127,7 +133,7 @@ class UniversalRelocBase } }); - auto address = resolveFunc(aHash); + auto address = resolveFunc(aSegment, aHash); if (address == 0) { std::stringstream stream; @@ -150,7 +156,7 @@ class UniversalRelocFunc : private UniversalRelocBase { public: UniversalRelocFunc(uint64_t aHash) - : m_address(reinterpret_cast(Resolve(aHash))) + : m_address(reinterpret_cast(Resolve(UniversalRelocSegment::Text, aHash))) { } @@ -171,8 +177,8 @@ template class UniversalRelocPtr : private UniversalRelocBase { public: - UniversalRelocPtr(uint64_t aHash) - : m_address(reinterpret_cast(Resolve(aHash))) + UniversalRelocPtr(uint64_t aHash, UniversalRelocSegment aSegment = UniversalRelocSegment::Data) + : m_address(reinterpret_cast(Resolve(aSegment, aHash))) { } @@ -198,7 +204,7 @@ class UniversalRelocVtbl : private UniversalRelocBase { public: UniversalRelocVtbl(uint64_t aHash) - : m_address(reinterpret_cast(Resolve(aHash))) + : m_address(reinterpret_cast(Resolve(UniversalRelocSegment::Rdata, aHash))) { }