1 files changed, 40 insertions, 2 deletions

diff --git a/plugins/msvc_rtti/rtti.cpp b/plugins/msvc_rtti/rtti.cpp
index fe89e739..ea129c39 100644
--- a/plugins/msvc_rtti/rtti.cpp
+++ b/plugins/msvc_rtti/rtti.cpp
@@ -569,13 +569,15 @@ std::optional<VirtualFunctionTableInfo> MicrosoftRTTIProcessor::ProcessVFT(uint6
 }
 
 
-MicrosoftRTTIProcessor::MicrosoftRTTIProcessor(const Ref<BinaryView> &view, bool useMangled, bool checkRData) : m_view(
+MicrosoftRTTIProcessor::MicrosoftRTTIProcessor(const Ref<BinaryView> &view, bool useMangled, bool checkRData,
+                                               bool vftSweep) : m_view(
     view)
 {
     m_logger = new Logger("Microsoft RTTI");
     allowMangledClassNames = useMangled;
     checkWritableRData = checkRData;
     m_classInfo = {};
+    virtualFunctionTableSweep = vftSweep;
     auto metadata = view->QueryMetadata(VIEW_METADATA_MSVC);
     if (metadata != nullptr)
     {
@@ -666,7 +668,43 @@ void MicrosoftRTTIProcessor::ProcessVFT()
         }
     }
 
+    if (virtualFunctionTableSweep)
+    {
+        BinaryReader optReader = BinaryReader(m_view);
+        auto addrSize = m_view->GetAddressSize();
+        auto scan = [&](const Ref<Segment> &segment) {
+            uint64_t startAddr = segment->GetStart();
+            uint64_t endAddr = segment->GetEnd();
+            for (uint64_t vtableAddr = startAddr; vtableAddr < endAddr - 0x18; vtableAddr += addrSize)
+            {
+                optReader.Seek(vtableAddr);
+                uint64_t coLocatorAddr = optReader.ReadPointer();
+                auto coLocator = m_classInfo.find(coLocatorAddr);
+                if (coLocator == m_classInfo.end())
+                    continue;
+                // Found a vtable reference to colocator.
+                ProcessVFT(vtableAddr + addrSize, coLocator->second);
+            }
+        };
+
+        // Scan data sections for virtual function tables.
+        auto rdataSection = m_view->GetSectionByName(".rdata");
+        for (const Ref<Segment> &segment: m_view->GetSegments())
+        {
+            if (segment->GetFlags() == (SegmentReadable | SegmentContainsData))
+            {
+                m_logger->LogDebug("Attempting to find VirtualFunctionTables in segment %llx", segment->GetStart());
+                scan(segment);
+            } else if (checkWritableRData && rdataSection && rdataSection->GetStart() == segment->GetStart())
+            {
+                m_logger->LogDebug("Attempting to find VirtualFunctionTables in writable rdata segment %llx",
+                                   segment->GetStart());
+                scan(segment);
+            }
+        }
+    }
+
     auto end_time = std::chrono::high_resolution_clock::now();
     std::chrono::duration<double> elapsed_time = end_time - start_time;
     m_logger->LogInfo("ProcessVFT took %f seconds", elapsed_time.count());
-}
\ No newline at end of file
+}