From c3eba484723475a11cb4df93c7906e959ecd25c2 Mon Sep 17 00:00:00 2001 From: Kawe Mazidjatari <48657826+Mauler125@users.noreply.github.com> Date: Fri, 2 Dec 2022 10:37:41 +0100 Subject: [PATCH] Implement signature cache further into the CModule class Init speeds have been reduced from 1.7 sec (average) to 0.06 sec (average) --- r5dev/public/utility/module.cpp | 49 +++++++++++++++++++++++-------- r5dev/public/utility/sigcache.cpp | 31 +++++++++++++++++++ r5dev/public/utility/sigcache.h | 1 + 3 files changed, 69 insertions(+), 12 deletions(-) diff --git a/r5dev/public/utility/module.cpp b/r5dev/public/utility/module.cpp index f98ebbd4..cd757698 100644 --- a/r5dev/public/utility/module.cpp +++ b/r5dev/public/utility/module.cpp @@ -118,21 +118,17 @@ CMemory CModule::FindPatternSIMD(const uint8_t* szPattern, const char* szMask, c //----------------------------------------------------------------------------- -// Purpose: find array of bytes in process memory using SIMD instructions -// Input : *svPattern +// Purpose: find a string pattern in process memory using SIMD instructions +// Input : &svPattern +// &moduleSection // Output : CMemory //----------------------------------------------------------------------------- CMemory CModule::FindPatternSIMD(const string& svPattern, const ModuleSections_t& moduleSection) const { - if (g_SigCache.m_bInitialized) // Get from cache instead. + uint64_t nRVA; + if (g_SigCache.FindEntry(svPattern, nRVA)) { - google::protobuf::Map sMap = g_SigCache.m_Cache.smap(); - - auto p = sMap.find(svPattern); - if (p != sMap.end()) - { - return CMemory((GetModuleBase() + p->second)); - } + return CMemory(nRVA + GetModuleBase()); } const pair patternInfo = PatternToMaskedBytes(svPattern); @@ -193,6 +189,14 @@ CMemory CModule::FindString(const string& svString, const ptrdiff_t nOccurrence, if (!m_ExecutableCode.IsSectionValid()) return CMemory(); + uint64_t nRVA; + string svPackedString = svString + std::to_string(nOccurrence); + + if (g_SigCache.FindEntry(svPackedString, nRVA)) + { + return CMemory(nRVA + GetModuleBase()); + } + const CMemory stringAddress = FindStringReadOnly(svString, bNullTerminator); // Get Address for the string in the .rdata section. if (!stringAddress) @@ -201,6 +205,7 @@ CMemory CModule::FindString(const string& svString, const ptrdiff_t nOccurrence, uint8_t* pLatestOccurrence = nullptr; uint8_t* pTextStart = reinterpret_cast(m_ExecutableCode.m_pSectionBase); // Get the start of the .text section. ptrdiff_t dOccurrencesFound = 0; + CMemory resultAddress; for (size_t i = 0ull; i < m_ExecutableCode.m_nSectionSize - 0x5; i++) { @@ -216,13 +221,22 @@ CMemory CModule::FindString(const string& svString, const ptrdiff_t nOccurrence, { dOccurrencesFound++; if (nOccurrence == dOccurrencesFound) - return CMemory(&pTextStart[i]); + { + resultAddress = CMemory(&pTextStart[i]); + g_SigCache.AddEntry(svPackedString, GetRVA(resultAddress.GetPtr())); + + return resultAddress; + } pLatestOccurrence = &pTextStart[i]; // Stash latest occurrence. } } } - return CMemory(pLatestOccurrence); + + resultAddress = CMemory(pLatestOccurrence); + + g_SigCache.AddEntry(svPackedString, GetRVA(resultAddress.GetPtr())); + return resultAddress; } //----------------------------------------------------------------------------- @@ -285,6 +299,14 @@ CMemory CModule::GetExportedFunction(const string& svFunctionName) const //----------------------------------------------------------------------------- CMemory CModule::GetVirtualMethodTable(const string& svTableName, const uint32_t nRefIndex) { + uint64_t nRVA; // Packed together as we can have multiple VFTable searches, but with different ref indexes. + string svPackedTableName = svTableName + std::to_string(nRefIndex); + + if (g_SigCache.FindEntry(svPackedTableName, nRVA)) + { + return CMemory(nRVA + GetModuleBase()); + } + const auto tableNameInfo = StringToMaskedBytes(svTableName, false); CMemory rttiTypeDescriptor = FindPatternSIMD(tableNameInfo.first.data(), tableNameInfo.second.c_str(), { ".data", m_RunTimeData.m_pSectionBase, m_RunTimeData.m_nSectionSize }).OffsetSelf(-0x10); if (!rttiTypeDescriptor) @@ -307,6 +329,9 @@ CMemory CModule::GetVirtualMethodTable(const string& svTableName, const uint32_t continue; } + CMemory vfTable = FindPatternSIMD(reinterpret_cast(&referenceOffset), "xxxxxxxx", { ".rdata", m_ReadOnlyData.m_pSectionBase, m_ReadOnlyData.m_nSectionSize }).OffsetSelf(0x8); + + g_SigCache.AddEntry(svPackedTableName, GetRVA(vfTable.GetPtr())); return FindPatternSIMD(reinterpret_cast(&referenceOffset), "xxxxxxxx", { ".rdata", m_ReadOnlyData.m_pSectionBase, m_ReadOnlyData.m_nSectionSize }).OffsetSelf(0x8); } diff --git a/r5dev/public/utility/sigcache.cpp b/r5dev/public/utility/sigcache.cpp index db7a9023..2eef7a54 100644 --- a/r5dev/public/utility/sigcache.cpp +++ b/r5dev/public/utility/sigcache.cpp @@ -7,6 +7,11 @@ #include "public/utility/binstream.h" #include "sigcache.h" +//----------------------------------------------------------------------------- +// Purpose: creates a pair of a pattern (key) and relative virtual address (value) +// Input : &svPattern - +// nRVA - +//----------------------------------------------------------------------------- void CSigCache::AddEntry(const string& svPattern, const uint64_t nRVA) { if (g_SigCache.m_bUseCache) @@ -15,6 +20,32 @@ void CSigCache::AddEntry(const string& svPattern, const uint64_t nRVA) } } +//----------------------------------------------------------------------------- +// Purpose: finds a pattern key in the cache map and sets its value to nRVA +// Input : &svPattern - +// &nRVA - +// Output : true if key is found, false otherwise +//----------------------------------------------------------------------------- +bool CSigCache::FindEntry(const string& svPattern, uint64_t& nRVA) const +{ + if (g_SigCache.m_bInitialized) + { + google::protobuf::Map sMap = g_SigCache.m_Cache.smap(); + auto p = sMap.find(svPattern); + + if (p != sMap.end()) + { + nRVA = p->second; + return true; + } + } + + return false; +} + +//----------------------------------------------------------------------------- +// Purpose: writes the cache map to the disk +//----------------------------------------------------------------------------- void CSigCache::WriteCache() { CIOStream writer("bin\\startup.smap", CIOStream::Mode_t::WRITE); diff --git a/r5dev/public/utility/sigcache.h b/r5dev/public/utility/sigcache.h index 2f3b7c7c..b658fc00 100644 --- a/r5dev/public/utility/sigcache.h +++ b/r5dev/public/utility/sigcache.h @@ -15,6 +15,7 @@ public: // Clear void AddEntry(const string& svPattern, const uint64_t nRVA); + bool FindEntry(const string& svPattern, uint64_t& nRVA) const; void WriteCache(); SigMap_Pb m_Cache;