RTTI: Add background task back and make it cancellable

If the background task gets stuck, this is the commit to ponder at

2 files changed, 16 insertions, 0 deletions

diff --git a/plugins/rtti/itanium.cpp b/plugins/rtti/itanium.cpp
index c3a2be5d..122869cc 100644
--- a/plugins/rtti/itanium.cpp
+++ b/plugins/rtti/itanium.cpp
@@ -659,6 +659,7 @@ ItaniumRTTIProcessor::ItaniumRTTIProcessor(const Ref<BinaryView> &view, bool use
 
 void ItaniumRTTIProcessor::ProcessRTTI()
 {
+    auto bgTask = new BackgroundTask("Scanning for Itanium RTTI...", true);
     auto start_time = std::chrono::high_resolution_clock::now();
     auto addrSize = m_view->GetAddressSize();
     uint64_t maxTypeInfoSize = TypeInfoSize(m_view);
@@ -666,6 +667,8 @@ void ItaniumRTTIProcessor::ProcessRTTI()
     auto scan = [&](const Ref<Section> &section) {
         for (uint64_t currAddr = section->GetStart(); currAddr <= section->GetEnd() - maxTypeInfoSize; currAddr += addrSize)
         {
+            if (bgTask->IsCancelled())
+                break;
             try
             {
                 if (auto classInfo = ProcessRTTI(currAddr))
@@ -732,6 +735,7 @@ void ItaniumRTTIProcessor::ProcessRTTI()
         );
     }
 
+    bgTask->Finish();
     auto end_time = std::chrono::high_resolution_clock::now();
     std::chrono::duration<double> elapsed_time = end_time - start_time;
     m_logger->LogDebug("ProcessRTTI took %f seconds", elapsed_time.count());
@@ -740,6 +744,7 @@ void ItaniumRTTIProcessor::ProcessRTTI()
 
 void ItaniumRTTIProcessor::ProcessVFT()
 {
+    auto bgTask = new BackgroundTask("Scanning for Itanium VFTs...", true);
     BinaryReader optReader = BinaryReader(m_view);
     std::map<uint64_t, std::set<uint64_t>> vftMap = {};
     std::map<uint64_t, std::optional<VirtualFunctionTableInfo>> vftFinishedMap = {};
@@ -811,12 +816,15 @@ void ItaniumRTTIProcessor::ProcessVFT()
 
     for (const auto &[coLocatorAddr, vftAddrs]: vftMap)
     {
+        if (bgTask->IsCancelled())
+            break;
         for (const auto& vftAddr: vftAddrs)
         {
             populateVftEntries(coLocatorAddr, vftAddr);
         }
     }
 
+    bgTask->Finish();
     auto end_time = std::chrono::high_resolution_clock::now();
     std::chrono::duration<double> elapsed_time = end_time - start_time;
     m_logger->LogDebug("ProcessVFT took %f seconds", elapsed_time.count());
diff --git a/plugins/rtti/microsoft.cpp b/plugins/rtti/microsoft.cpp
index 66015042..95c76311 100644
--- a/plugins/rtti/microsoft.cpp
+++ b/plugins/rtti/microsoft.cpp
@@ -620,6 +620,7 @@ MicrosoftRTTIProcessor::MicrosoftRTTIProcessor(const Ref<BinaryView> &view, bool
 
 void MicrosoftRTTIProcessor::ProcessRTTI()
 {
+    auto bgTask = new BackgroundTask("Scanning for Microsoft RTTI...", true);
     auto start_time = std::chrono::high_resolution_clock::now();
     uint64_t startAddr = m_view->GetOriginalImageBase();
     uint64_t endAddr = m_view->GetEnd();
@@ -630,6 +631,8 @@ void MicrosoftRTTIProcessor::ProcessRTTI()
         for (uint64_t coLocatorAddr = segment->GetStart(); coLocatorAddr < segment->GetEnd() - 0x18;
              coLocatorAddr += addrSize)
         {
+            if (bgTask->IsCancelled())
+                break;
             optReader.Seek(coLocatorAddr);
             uint32_t sigVal = optReader.Read32();
             if (sigVal == COL_SIG_REV1)
@@ -683,6 +686,7 @@ void MicrosoftRTTIProcessor::ProcessRTTI()
         }
     }
 
+    bgTask->Finish();
     auto end_time = std::chrono::high_resolution_clock::now();
     std::chrono::duration<double> elapsed_time = end_time - start_time;
     m_logger->LogDebug("ProcessRTTI took %f seconds", elapsed_time.count());
@@ -691,6 +695,7 @@ void MicrosoftRTTIProcessor::ProcessRTTI()
 
 void MicrosoftRTTIProcessor::ProcessVFT()
 {
+    auto bgTask = new BackgroundTask("Scanning for Microsoft RTTI...", true);
     std::map<uint64_t, uint64_t> vftMap = {};
     std::map<uint64_t, std::optional<VirtualFunctionTableInfo>> vftFinishedMap = {};
     auto start_time = std::chrono::high_resolution_clock::now();
@@ -712,6 +717,8 @@ void MicrosoftRTTIProcessor::ProcessVFT()
             uint64_t endAddr = segment->GetEnd();
             for (uint64_t vtableAddr = startAddr; vtableAddr < endAddr - 0x18; vtableAddr += addrSize)
             {
+                if (bgTask->IsCancelled())
+                    break;
                 optReader.Seek(vtableAddr);
                 uint64_t coLocatorAddr = optReader.ReadPointer();
                 auto coLocator = m_classInfo.find(coLocatorAddr);
@@ -806,6 +813,7 @@ void MicrosoftRTTIProcessor::ProcessVFT()
     for (const auto &[coLocatorAddr, _]: vftMap)
         ProcessClassAndBases(coLocatorAddr);
 
+    bgTask->Finish();
     auto end_time = std::chrono::high_resolution_clock::now();
     std::chrono::duration<double> elapsed_time = end_time - start_time;
     m_logger->LogDebug("ProcessVFT took %f seconds", elapsed_time.count());