MSVC RTTI: Recursively define ClassHierarchyDescriptors

This will better define base classes without a CompleteObjectLocator
associated with it

2 files changed, 72 insertions, 33 deletions

diff --git a/plugins/msvc_rtti/rtti.cpp b/plugins/msvc_rtti/rtti.cpp
index afaa8424..f1d172a1 100644
--- a/plugins/msvc_rtti/rtti.cpp
+++ b/plugins/msvc_rtti/rtti.cpp
@@ -452,6 +452,64 @@ std::optional<ClassInfo> MicrosoftRTTIProcessor::ProcessRTTI(uint64_t coLocatorA
 
     auto ptrBaseTy = coLocator->signature ? RelativeToBinaryStartPointerBaseType : AbsolutePointerBaseType;
 
+    auto defineClassHierarchyDesc = [&](const uint64_t classHierarchyDescAddr, ClassInfo& classInfo, std::optional<CompleteObjectLocator> coLocator) {
+        auto classHierarchyDesc = ClassHierarchyDescriptor(m_view, classHierarchyDescAddr);
+        auto classHierarchyDescName = fmt::format("{}::`RTTI Class Hierarchy Descriptor'", classInfo.className);
+        m_view->DefineAutoSymbol(new Symbol{DataSymbol, classHierarchyDescName, classHierarchyDescAddr});
+        m_view->DefineDataVariable(classHierarchyDescAddr,
+                                Confidence(ClassHierarchyDescriptorType(m_view, ptrBaseTy), RTTI_CONFIDENCE));
+
+        auto baseClassArrayAddr = resolveAddr(classHierarchyDesc.pBaseClassArray);
+        auto baseClassArray = BaseClassArray(m_view, baseClassArrayAddr, classHierarchyDesc.numBaseClasses);
+        auto baseClassArrayName = fmt::format("{}::`RTTI Base Class Array'", classInfo.className);
+        m_view->DefineAutoSymbol(new Symbol{DataSymbol, baseClassArrayName, baseClassArrayAddr});
+        m_view->DefineDataVariable(baseClassArrayAddr,
+                                Confidence(BaseClassArrayType(m_view, baseClassArray.length, ptrBaseTy),
+                                            RTTI_CONFIDENCE));
+
+        std::map<uint64_t, ClassInfo> baseClasses = {};
+        for (auto pBaseClassDescAddr: baseClassArray.descriptors)
+        {
+            auto baseClassDescAddr = resolveAddr(pBaseClassDescAddr);
+            auto baseClassDesc = BaseClassDescriptor(m_view, baseClassDescAddr);
+
+            auto baseClassTypeDescAddr = resolveAddr(baseClassDesc.pTypeDescriptor);
+            auto baseClassTypeDesc = TypeDescriptor(m_view, baseClassTypeDescAddr);
+            auto baseClassName = DemangleName(baseClassTypeDesc.name);
+            if (!baseClassName.has_value())
+            {
+                m_logger->LogWarn("Skipping BaseClassDescriptor with mangled name %llx", baseClassTypeDescAddr);
+                continue;
+            }
+
+            // TODO: we probably want to maintain this state
+            auto baseClassInfo = ClassInfo{baseClassName.value()};
+
+            if (coLocator.has_value())
+            {
+                if (baseClassDesc.where_mdisp == coLocator->offset && !classInfo.baseClassName.has_value() && classInfo.className != baseClassInfo.className)
+                    classInfo.baseClassName = baseClassInfo.className;
+            }
+
+            auto baseClassDescName = fmt::format("{}::`RTTI Base Class Descriptor at ({},{},{},{})", baseClassInfo.className,
+                                                baseClassDesc.where_mdisp, baseClassDesc.where_pdisp,
+                                                baseClassDesc.where_vdisp, baseClassDesc.attributes);
+            m_view->DefineAutoSymbol(new Symbol{DataSymbol, baseClassDescName, baseClassDescAddr});
+            m_view->DefineDataVariable(baseClassDescAddr,
+                                    Confidence(BaseClassDescriptorType(m_view, ptrBaseTy), RTTI_CONFIDENCE));
+
+            auto baseClassTypeDescSymName = fmt::format("class {} `RTTI Type Descriptor'", baseClassInfo.className);
+            m_view->DefineAutoSymbol(new Symbol{DataSymbol, baseClassTypeDescSymName, baseClassTypeDescAddr});
+            m_view->DefineDataVariable(baseClassTypeDescAddr,
+                                    Confidence(TypeDescriptorType(m_view, baseClassTypeDesc.name.length()), RTTI_CONFIDENCE));
+
+            auto classHierarchyDescAddr = resolveAddr(baseClassDesc.pClassHierarchyDescriptor);
+            baseClasses[classHierarchyDescAddr] = baseClassInfo;
+        }
+
+        return baseClasses;
+    };
+
     // Get type descriptor then check to see if the class name was demangled.
     auto typeDescAddr = resolveAddr(coLocator->pTypeDescriptor);
     auto typeDesc = TypeDescriptor(m_view, typeDescAddr);
@@ -480,43 +538,21 @@ std::optional<ClassInfo> MicrosoftRTTIProcessor::ProcessRTTI(uint64_t coLocatorA
                                Confidence(TypeDescriptorType(m_view, typeDesc.name.length()), RTTI_CONFIDENCE));
 
     auto classHierarchyDescAddr = resolveAddr(coLocator->pClassHierarchyDescriptor);
-    auto classHierarchyDesc = ClassHierarchyDescriptor(m_view, classHierarchyDescAddr);
-    auto classHierarchyDescName = fmt::format("{}::`RTTI Class Hierarchy Descriptor'", classInfo.className);
-    m_view->DefineAutoSymbol(new Symbol{DataSymbol, classHierarchyDescName, classHierarchyDescAddr});
-    m_view->DefineDataVariable(classHierarchyDescAddr,
-                               Confidence(ClassHierarchyDescriptorType(m_view, ptrBaseTy), RTTI_CONFIDENCE));
-
-    auto baseClassArrayAddr = resolveAddr(classHierarchyDesc.pBaseClassArray);
-    auto baseClassArray = BaseClassArray(m_view, baseClassArrayAddr, classHierarchyDesc.numBaseClasses);
-    auto baseClassArrayName = fmt::format("{}::`RTTI Base Class Array'", classInfo.className);
-    m_view->DefineAutoSymbol(new Symbol{DataSymbol, baseClassArrayName, baseClassArrayAddr});
-    m_view->DefineDataVariable(baseClassArrayAddr,
-                               Confidence(BaseClassArrayType(m_view, baseClassArray.length, ptrBaseTy),
-                                          RTTI_CONFIDENCE));
-
-    for (auto pBaseClassDescAddr: baseClassArray.descriptors)
+    auto baseClasses = defineClassHierarchyDesc(classHierarchyDescAddr, classInfo, coLocator);
+    m_visitedClassHierarchyDescAddrs.insert(classHierarchyDescAddr);
+    while (baseClasses.size() > 0)
     {
-        auto baseClassDescAddr = resolveAddr(pBaseClassDescAddr);
-        auto baseClassDesc = BaseClassDescriptor(m_view, baseClassDescAddr);
-
-        auto baseClassTypeDescAddr = resolveAddr(baseClassDesc.pTypeDescriptor);
-        auto baseClassTypeDesc = TypeDescriptor(m_view, baseClassTypeDescAddr);
-        auto baseClassName = DemangleName(baseClassTypeDesc.name);
-        if (!baseClassName.has_value())
+        std::map<uint64_t, ClassInfo> newBaseClasses = {};
+        for (auto& [baseClassHierarchyDescAddr, baseClassInfo] : baseClasses)
         {
-            m_logger->LogWarn("Skipping BaseClassDescriptor with mangled name %llx", baseClassTypeDescAddr);
-            continue;
-        }
+            if (m_visitedClassHierarchyDescAddrs.find(baseClassHierarchyDescAddr) != m_visitedClassHierarchyDescAddrs.end())
+                continue;
 
-        if (baseClassDesc.where_mdisp == coLocator->offset && classInfo.className != baseClassName.value())
-            classInfo.baseClassName = baseClassName;
+            newBaseClasses.merge(defineClassHierarchyDesc(baseClassHierarchyDescAddr, baseClassInfo, std::nullopt));
+            m_visitedClassHierarchyDescAddrs.insert(baseClassHierarchyDescAddr);
+        }
 
-        auto baseClassDescName = fmt::format("{}::`RTTI Base Class Descriptor at ({},{},{},{})", baseClassName.value(),
-                                             baseClassDesc.where_mdisp, baseClassDesc.where_pdisp,
-                                             baseClassDesc.where_vdisp, baseClassDesc.attributes);
-        m_view->DefineAutoSymbol(new Symbol{DataSymbol, baseClassDescName, baseClassDescAddr});
-        m_view->DefineDataVariable(baseClassDescAddr,
-                                   Confidence(BaseClassDescriptorType(m_view, ptrBaseTy), RTTI_CONFIDENCE));
+        baseClasses = newBaseClasses;
     }
 
     auto coLocatorName = fmt::format("{}::`RTTI Complete Object Locator'", className.value());
@@ -661,6 +697,7 @@ MicrosoftRTTIProcessor::MicrosoftRTTIProcessor(const Ref<BinaryView> &view, bool
     allowAnonymousClassNames = allowAnonymous;
     checkWritableRData = checkRData;
     m_classInfo = {};
+    m_visitedClassHierarchyDescAddrs = {};
     virtualFunctionTableSweep = vftSweep;
     auto metadata = view->QueryMetadata(VIEW_METADATA_MSVC);
     if (metadata != nullptr)
diff --git a/plugins/msvc_rtti/rtti.h b/plugins/msvc_rtti/rtti.h
index ab964ba9..0bb3a733 100644
--- a/plugins/msvc_rtti/rtti.h
+++ b/plugins/msvc_rtti/rtti.h
@@ -101,6 +101,8 @@ namespace BinaryNinja {
 
 		std::map<uint64_t, ClassInfo> m_classInfo;
 
+		std::set<uint64_t> m_visitedClassHierarchyDescAddrs;
+
 		void DeserializedMetadata(const Ref<Metadata> &metadata);
 
 		std::optional<std::string> DemangleName(const std::string &mangledName);