Add segment type to the resolve function

wopss · Feb 8, 2024 · 9d54397 · 9d54397
1 parent b0e1d9f
commit 9d54397
Showing 1 changed file with 15 additions and 9 deletions.
diff --git 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<uintptr_t (*)(uint64_t)>(GetProcAddress(handle, functionName));
+                           resolveFunc = reinterpret_cast<functionType>(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<T>(Resolve(aHash)))
+        : m_address(reinterpret_cast<T>(Resolve(UniversalRelocSegment::Text, aHash)))
     {
     }
 
@@ -171,8 +177,8 @@ template<typename T>
 class UniversalRelocPtr : private UniversalRelocBase
 {
 public:
-    UniversalRelocPtr(uint64_t aHash)
-        : m_address(reinterpret_cast<T*>(Resolve(aHash)))
+    UniversalRelocPtr(uint64_t aHash, UniversalRelocSegment aSegment = UniversalRelocSegment::Data)
+        : m_address(reinterpret_cast<T*>(Resolve(aSegment, aHash)))
     {
     }
 
@@ -198,7 +204,7 @@ class UniversalRelocVtbl : private UniversalRelocBase
 {
 public:
     UniversalRelocVtbl(uint64_t aHash)
-        : m_address(reinterpret_cast<uintptr_t*>(Resolve(aHash)))
+        : m_address(reinterpret_cast<uintptr_t*>(Resolve(UniversalRelocSegment::Rdata, aHash)))
     {
     }