fix: prevent VirtualFunction memory leak in native layer (#1121)

Author: ELDment

src/core/function.cpp

@@ -98,7 +98,20 @@ ValveFunction::ValveFunction(void* ulAddr, Convention_t callingConvention, DataT
     m_iCallingConvention = GetDynCallConvention(m_eCallingConvention);
 }
 
-ValveFunction::~ValveFunction() {}
+ValveFunction::~ValveFunction()
+{
+    if (m_precallback != nullptr)
+    {
+        globals::callbackManager.ReleaseCallback(m_precallback);
+        m_precallback = nullptr;
+    }
+
+    if (m_postcallback != nullptr)
+    {
+        globals::callbackManager.ReleaseCallback(m_postcallback);
+        m_postcallback = nullptr;
+    }
+}
 
 bool ValveFunction::IsCallable() { return (m_eCallingConvention != CONV_CUSTOM) && (m_iCallingConvention != -1); }
 

src/scripting/natives/natives_memory.cpp

@@ -16,6 +16,7 @@
 
 #include <ios>
 #include <sstream>
+#include <unordered_map>
 
 #include "core/function.h"
 #include "core/log.h"
@@ -24,7 +25,51 @@
 #include "scripting/script_engine.h"
 
 namespace counterstrikesharp {
-std::vector<ValveFunction*> m_managed_ptrs;
+
+template <class T> inline void hash_combine(std::size_t& seed, const T& v)
+{
+    std::hash<T> hasher;
+    seed ^= hasher(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+}
+
+struct VirtualFunctionCacheKey
+{
+    void* functionAddr;
+    Convention_t callingConvention;
+    std::vector<DataType_t> args;
+    DataType_t returnType;
+    int vtableOffset;
+
+    bool operator==(const VirtualFunctionCacheKey& other) const
+    {
+        return functionAddr == other.functionAddr && callingConvention == other.callingConvention && args == other.args &&
+               returnType == other.returnType && vtableOffset == other.vtableOffset;
+    }
+};
+
+struct VirtualFunctionCacheKeyHash
+{
+    std::size_t operator()(const VirtualFunctionCacheKey& key) const
+    {
+        std::size_t hash = 0;
+
+        hash_combine(hash, std::hash<void*>{}(key.functionAddr));
+        hash_combine(hash, std::hash<int>{}(static_cast<int>(key.callingConvention)));
+        hash_combine(hash, std::hash<int>{}(static_cast<int>(key.returnType)));
+        hash_combine(hash, std::hash<int>{}(key.vtableOffset));
+
+        for (const auto& arg : key.args)
+        {
+            hash_combine(hash, std::hash<int>{}(static_cast<int>(arg)));
+        }
+
+        return hash;
+    }
+};
+
+std::unordered_map<VirtualFunctionCacheKey, ValveFunction*, VirtualFunctionCacheKeyHash> m_virtualFunctionCache;
+
+size_t GetVirtualFunctionCacheSize() { return m_virtualFunctionCache.size(); }
 
 void* FindSignatureNative(ScriptContext& scriptContext)
 {
@@ -64,12 +109,28 @@ ValveFunction* CreateVirtualFunctionBySignature(ScriptContext& script_context)
         args.push_back(script_context.GetArgument<DataType_t>(5 + i));
     }
 
+    VirtualFunctionCacheKey cacheKey;
+    cacheKey.functionAddr = function_addr;
+    cacheKey.callingConvention = CONV_CDECL;
+    cacheKey.args = args;
+    cacheKey.returnType = return_type;
+    cacheKey.vtableOffset = -1;
+
+    auto it = m_virtualFunctionCache.find(cacheKey);
+    if (it != m_virtualFunctionCache.end())
+    {
+        CSSHARP_CORE_TRACE("Virtual function found in cache, reusing existing instance at {}, signature {}", function_addr,
+                           signature_hex_string);
+        return it->second;
+    }
+
     auto function = new ValveFunction(function_addr, CONV_CDECL, args, return_type);
     function->SetSignature(signature_hex_string);
 
-    CSSHARP_CORE_TRACE("Created virtual function, pointer found at {}, signature {}", function_addr, signature_hex_string);
+    m_virtualFunctionCache[cacheKey] = function;
+
+    CSSHARP_CORE_TRACE("Created new virtual function, pointer found at {}, signature {}", function_addr, signature_hex_string);
 
-    m_managed_ptrs.push_back(function);
     return function;
 }
 
@@ -95,10 +156,27 @@ ValveFunction* CreateVirtualFunction(ScriptContext& script_context)
         args.push_back(script_context.GetArgument<DataType_t>(4 + i));
     }
 
+    VirtualFunctionCacheKey cacheKey;
+    cacheKey.functionAddr = function_addr;
+    cacheKey.callingConvention = CONV_THISCALL;
+    cacheKey.args = args;
+    cacheKey.returnType = return_type;
+    cacheKey.vtableOffset = vtable_offset;
+
+    auto it = m_virtualFunctionCache.find(cacheKey);
+    if (it != m_virtualFunctionCache.end())
+    {
+        CSSHARP_CORE_TRACE("Virtual function found in cache, reusing existing instance at {}, offset {}", function_addr, vtable_offset);
+        return it->second;
+    }
+
     auto function = new ValveFunction(function_addr, CONV_THISCALL, args, return_type);
     function->SetOffset(vtable_offset);
 
-    m_managed_ptrs.push_back(function);
+    m_virtualFunctionCache[cacheKey] = function;
+
+    CSSHARP_CORE_TRACE("Created new virtual function at {}, offset {}", function_addr, vtable_offset);
+
     return function;
 }
 

src/scripting/natives/natives_vector.cpp

@@ -41,17 +41,19 @@ CREATE_SETTER_FUNCTION(Vector, float, Z, Vector*, obj->z = value);
 std::vector<Vector*> managed_vectors;
 std::vector<QAngle*> managed_angles;
 extern std::vector<IGameEvent*> managed_game_events;
-extern std::vector<ValveFunction*> m_managed_ptrs;
+
+extern size_t GetVirtualFunctionCacheSize();
 
 CON_COMMAND(css_dump_leaks, "dump css leaks")
 {
+    auto virtualFunctionCount = GetVirtualFunctionCacheSize();
     Msg("===== Dumping leaks =====\n");
     Msg("\tVector: %i (%zu B)\n", managed_vectors.size(), managed_vectors.size() * sizeof(Vector));
     Msg("\tAngles: %i (%zu B)\n", managed_angles.size(), managed_angles.size() * sizeof(QAngle));
     Msg("\tGameEvents: %i (~B)\n", managed_game_events.size());
-    Msg("\tVirtual Functions: %i (%zu B)\n", m_managed_ptrs.size(), m_managed_ptrs.size() * sizeof(ValveFunction));
+    Msg("\tVirtual Functions: %i (%zu B)\n", virtualFunctionCount, virtualFunctionCount * sizeof(ValveFunction));
     Msg("\tTotal size: %zu B\n", (managed_vectors.size() * sizeof(Vector)) + (managed_angles.size() * sizeof(QAngle)) +
-                                     (m_managed_ptrs.size() * sizeof(ValveFunction)));
+                                     (virtualFunctionCount * sizeof(ValveFunction)));
     Msg("===== Dumping leaks =====\n");
 }