From 40870d62ab47233b83ef1496f23bf349869443d2 Mon Sep 17 00:00:00 2001
From: praydog <praydog@praydog.com>
Date: Sun, 22 Oct 2023 07:52:08 -0700
Subject: [PATCH] SDK: Add scanning of Malloc, Realloc, Free indices

---
 shared/sdk/FMalloc.cpp | 86 ++++++++++++++++++++++++++++++++++++++++++
 shared/sdk/FMalloc.hpp | 27 +++++++++++++
 2 files changed, 113 insertions(+)

diff --git a/shared/sdk/FMalloc.cpp b/shared/sdk/FMalloc.cpp
index bc1b8398..b27e6660 100644
--- a/shared/sdk/FMalloc.cpp
+++ b/shared/sdk/FMalloc.cpp
@@ -122,6 +122,92 @@ FMalloc* FMalloc::get() {
             return nullptr;
         }
 
+        // Since we found it now, go through the vtable and try to determine the indices
+        // for Malloc, Realloc, and Free.
+        const auto vtable = *(void***)*result;
+
+        for (auto i = 1; i < 30; ++i) {
+            const auto fn = (uintptr_t)vtable[i];
+
+            if (IsBadReadPtr((void*)fn, sizeof(void*))) {
+                break;
+            }
+
+            if (utility::find_pointer_in_path(fn, &VirtualAlloc)) {
+                SPDLOG_INFO("[FMalloc::get] Found Malloc at index {}", i);
+                s_malloc_index = i;
+                break;
+            }
+        }
+
+        if (!s_malloc_index.has_value()) {
+            SPDLOG_ERROR("[FMalloc::get] Failed to find FMalloc::Malloc");
+            return result;
+        }
+
+        /*for (auto i = *s_malloc_index + 1; i < 30; ++i) {
+            const auto fn = (uintptr_t)vtable[i];
+
+            if (IsBadReadPtr((void*)fn, sizeof(void*))) {
+                break;
+            }
+
+            if (utility::find_pointer_in_path(fn, &VirtualAlloc) || utility::find_pointer_in_path(fn, &VirtualFree)) {
+                SPDLOG_INFO("[FMalloc::get] Found Realloc at index {}", i);
+                s_realloc_index = i;
+                break;
+            }
+        }*/
+
+        // The right one is whichever one doesn't return really early after malloc.
+        // We can scan for VirtualAlloc/Free in SOME builds, but this doesn't always work.
+        // Because in some builds, Realloc calls Malloc via an indirect call, which exhaustive_decode doesn't handle.
+        // I could handle it with emulation but I don't want to.
+        for (auto i = *s_malloc_index + 1; i < 30; ++i) {
+            const auto fn = (uintptr_t)vtable[i];
+
+            if (IsBadReadPtr((void*)fn, sizeof(void*))) {
+                break;
+            }
+
+            uint32_t distance_from_ret = 0;
+
+            utility::exhaustive_decode((uint8_t*)fn, 100, [&](utility::ExhaustionContext& ctx) -> utility::ExhaustionResult {
+                ++distance_from_ret;
+                return utility::ExhaustionResult::CONTINUE;
+            });
+
+            if (distance_from_ret > 10) {
+                s_realloc_index = i;
+                SPDLOG_INFO("[FMalloc::get] Found Realloc at index {}", i);
+                break;
+            }
+        }
+
+        if (!s_realloc_index.has_value()) {
+            SPDLOG_ERROR("[FMalloc::get] Failed to find FMalloc::Realloc");
+            return result;
+        }
+
+        for (auto i = *s_realloc_index + 1; i < 30; ++i) {
+            const auto fn = (uintptr_t)vtable[i];
+
+            if (IsBadReadPtr((void*)fn, sizeof(void*))) {
+                break;
+            }
+
+            if (utility::find_pointer_in_path(fn, &VirtualFree)) {
+                SPDLOG_INFO("[FMalloc::get] Found Free at index {}", i);
+                s_free_index = i;
+                break;
+            }
+        }
+
+        if (!s_free_index.has_value()) {
+            SPDLOG_ERROR("[FMalloc::get] Failed to find FMalloc::Free");
+            return result;
+        }
+
         return result;
     }();
 
diff --git a/shared/sdk/FMalloc.hpp b/shared/sdk/FMalloc.hpp
index 44c1a105..7d136ea0 100644
--- a/shared/sdk/FMalloc.hpp
+++ b/shared/sdk/FMalloc.hpp
@@ -7,6 +7,33 @@ class FMalloc {
 
     virtual ~FMalloc() {}
 
+    void* malloc(size_t size, uint32_t alignment = 0) {
+        if (!s_malloc_index.has_value()) {
+            return nullptr;
+        }
+
+        return ((void*(__thiscall*)(FMalloc*, size_t, uint32_t))(*(void***)this)[*s_malloc_index])(this, size, alignment);
+    }
+
+    void* realloc(void* original, size_t size, uint32_t alignment = 0) {
+        if (!s_realloc_index.has_value()) {
+            return nullptr;
+        }
+
+        return ((void*(__thiscall*)(FMalloc*, void*, size_t, uint32_t))(*(void***)this)[*s_realloc_index])(this, original, size, alignment);
+    }
+
+    void free(void* original) {
+        if (!s_free_index.has_value()) {
+            return;
+        }
+
+        ((void(__thiscall*)(FMalloc*, void*))(*(void***)this)[*s_free_index])(this, original);
+    }
+
 private:
+    static inline std::optional<uint32_t> s_malloc_index{};
+    static inline std::optional<uint32_t> s_realloc_index{};
+    static inline std::optional<uint32_t> s_free_index{};
 };
 }
\ No newline at end of file