From 81fd6aa419465d27b8109e6c658fb69a83e8b7bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Krupi=C5=84ski?= <danielkrupinski@outlook.com>
Date: Thu, 13 Jul 2023 22:20:49 +0200
Subject: [PATCH] Add getVmtSection() method to dynamic library classes

---
 Source/Platform/Linux/LinuxDynamicLibrary.h     | 14 ++++++++++++--
 Source/Platform/Windows/PortableExecutable.h    |  9 +++++++++
 Source/Platform/Windows/WindowsDynamicLibrary.h |  7 +++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/Source/Platform/Linux/LinuxDynamicLibrary.h b/Source/Platform/Linux/LinuxDynamicLibrary.h
index 9eb7e4e5e10..bffea9164b7 100644
--- a/Source/Platform/Linux/LinuxDynamicLibrary.h
+++ b/Source/Platform/Linux/LinuxDynamicLibrary.h
@@ -41,6 +41,17 @@ class LinuxDynamicLibrary {
     }
 
     [[nodiscard]] MemorySection getCodeSection() const noexcept
+    {
+        return getSection(".text");
+    }
+
+    [[nodiscard]] MemorySection getVmtSection() const noexcept
+    {
+        return getSection(".data.rel.ro");
+    }
+
+private:
+    [[nodiscard]] MemorySection getSection(const char* sectionName) const noexcept
     {
         void* base = nullptr;
         std::size_t size = 0;
@@ -57,7 +68,7 @@ class LinuxDynamicLibrary {
                         for (auto i = 0; i < ehdr->e_shnum; ++i) {
                             const auto shdr = (ElfW(Shdr)*)(std::uintptr_t(shdrs) + i * ehdr->e_shentsize);
 
-                            if (std::strcmp(strTab + shdr->sh_name, ".text") != 0)
+                            if (std::strcmp(strTab + shdr->sh_name, sectionName) != 0)
                                 continue;
 
                             base = (void*)(linkMap->l_addr + shdr->sh_offset);
@@ -75,7 +86,6 @@ class LinuxDynamicLibrary {
         return MemorySection{ std::span{ reinterpret_cast<const std::byte*>(base), size } };
     }
 
-private:
     [[nodiscard]] void* getModuleHandle(const char* libraryName)
     {
         const auto handle = PlatformApi::dlopen(libraryName, RTLD_LAZY | RTLD_NOLOAD);
diff --git a/Source/Platform/Windows/PortableExecutable.h b/Source/Platform/Windows/PortableExecutable.h
index 0e7b525caa5..c0e129062c2 100644
--- a/Source/Platform/Windows/PortableExecutable.h
+++ b/Source/Platform/Windows/PortableExecutable.h
@@ -27,6 +27,15 @@ class PortableExecutable {
         return {};
     }
 
+    [[nodiscard]] MemorySection getVmtSection() const noexcept
+    {
+        for (const auto& section : getSectionHeaders()) {
+            if ((section.Characteristics & IMAGE_SCN_MEM_READ) != 0 && std::memcmp(section.Name, ".rdata", 6) == 0)
+                return MemorySection{ std::span{ base + section.VirtualAddress, section.Misc.VirtualSize } };
+        }
+        return {};
+    }
+
     [[nodiscard]] SafeAddress getExport(const char* name) const noexcept
     {
         const auto exportDataDirectory = getExportDataDirectory();
diff --git a/Source/Platform/Windows/WindowsDynamicLibrary.h b/Source/Platform/Windows/WindowsDynamicLibrary.h
index a7260892afc..337633292c3 100644
--- a/Source/Platform/Windows/WindowsDynamicLibrary.h
+++ b/Source/Platform/Windows/WindowsDynamicLibrary.h
@@ -34,6 +34,13 @@ class WindowsDynamicLibrary {
         return {};
     }
 
+    [[nodiscard]] MemorySection getVmtSection() const noexcept
+    {
+        if (handle)
+            return portableExecutable().getVmtSection();
+        return {};
+    }
+
     [[nodiscard]] HMODULE getHandle() const noexcept
     {
         return handle;