[RTTI] Improve virtual function discovery

- Allow extern functions to show up in a MSVC vtable
- Fix https://github.com/Vector35/binaryninja-api/issues/7871
- Share code between itanium and msvc vft analysis, it is the same logic basically

4 files changed, 97 insertions, 80 deletions

diff --git a/plugins/rtti/itanium.cpp b/plugins/rtti/itanium.cpp
index 6f4c7158..145969c0 100644
--- a/plugins/rtti/itanium.cpp
+++ b/plugins/rtti/itanium.cpp
@@ -14,22 +14,6 @@ constexpr const char *TYPE_SOURCE_ITANIUM = "rtti_itanium";
 constexpr int MAX_FAILED_SCAN_ATTEMPTS = 10;
 
 
-Ref<Symbol> GetRealSymbol(BinaryView *view, uint64_t relocAddr, uint64_t symAddr)
-{
-    if (view->IsOffsetExternSemantics(symAddr))
-    {
-        // Because bases in the extern section are not 8 byte width only they will
-        // overlap with other externs, until https://github.com/Vector35/binaryninja-api/issues/6387 is fixed.
-        // Check relocation at objectAddr for symbol
-        for (const auto& r : view->GetRelocationsAt(relocAddr))
-            if (auto relocSym = r->GetSymbol())
-                return relocSym;
-    }
-
-    return view->GetSymbolByAddress(symAddr);
-}
-
-
 // Some fields are not always u32, use this if it goes from u32 -> u16 on 32bit
 uint64_t ArchFieldSize(BinaryView *view)
 {
@@ -37,7 +21,6 @@ uint64_t ArchFieldSize(BinaryView *view)
 }
 
 
-
 uint64_t TypeInfoSize(BinaryView *view)
 {
     return view->GetAddressSize() * 2;
@@ -558,47 +541,16 @@ std::optional<ClassInfo> ItaniumRTTIProcessor::ProcessRTTI(uint64_t objectAddr)
 std::optional<VirtualFunctionTableInfo> ItaniumRTTIProcessor::ProcessVFT(uint64_t vftAddr, ClassInfo &classInfo, std::optional<BaseClassInfo> baseClassInfo)
 {
     VirtualFunctionTableInfo vftInfo = {vftAddr};
-    BinaryReader reader = BinaryReader(m_view);
-    reader.Seek(vftAddr);
     // Gather all virtual functions
     std::vector<VirtualFunctionInfo> virtualFunctions = {};
+    uint64_t currentVftEntry = vftAddr;
     while (true)
     {
-        uint64_t readOffset = reader.GetOffset();
-        if (!m_view->IsValidOffset(readOffset))
+        uint64_t vFuncAddr = 0;
+        const FunctionDiscoverState state = DiscoverVirtualFunction(currentVftEntry, vFuncAddr);
+        if (state == FunctionDiscoverState::Failed)
             break;
-        uint64_t vFuncAddr = reader.ReadPointer();
-        auto funcs = m_view->GetAnalysisFunctionsForAddress(vFuncAddr);
-        if (funcs.empty())
-        {
-            Ref<Segment> segment = m_view->GetSegmentAt(vFuncAddr);
-            if (segment == nullptr || !(segment->GetFlags() & (SegmentExecutable | SegmentDenyWrite)))
-            {
-                // TODO: Sometimes vFunc idx will be zeroed iirc.
-                // We allow vfuncs to point to extern functions.
-                // TODO: Until https://github.com/Vector35/binaryninja-api/issues/5982 is fixed we need to check extern sym relocs instead of the symbol directly
-                auto vFuncSym = GetRealSymbol(m_view, reader.GetOffset(), vFuncAddr);
-                if (!vFuncSym)
-                    break;
-                DataVariable dv;
-                bool foundDv = m_view->GetDataVariableAtAddress(vFuncAddr, dv);
-                // Last virtual function, or hit the next vtable.
-                if (!foundDv || !dv.type->m_object)
-                    break;
-                // Void externs are very likely to be a func.
-                // TODO: Add some sanity checks for this!
-                if (!dv.type->IsFunction() && !(dv.type->IsVoid() && vFuncSym->GetType() == ExternalSymbol))
-                    break;
-            }
-            else
-            {
-                // TODO: Is likely a function check here?
-                m_logger->LogDebugF("Discovered function from virtual function table... {:#x}", vFuncAddr);
-                auto vftPlatform = m_view->GetDefaultPlatform()->GetAssociatedPlatformByAddress(vFuncAddr);
-                m_view->AddFunctionForAnalysis(vftPlatform, vFuncAddr, true);
-            }
-        }
-        // Only ever add one function.
+        currentVftEntry += m_view->GetAddressSize();
         virtualFunctions.emplace_back(VirtualFunctionInfo{vFuncAddr});
     }
 
diff --git a/plugins/rtti/microsoft.cpp b/plugins/rtti/microsoft.cpp
index 684a68cd..b57b15f6 100644
--- a/plugins/rtti/microsoft.cpp
+++ b/plugins/rtti/microsoft.cpp
@@ -482,37 +482,19 @@ std::optional<ClassInfo> MicrosoftRTTIProcessor::ProcessRTTI(uint64_t coLocatorA
 std::optional<VirtualFunctionTableInfo> MicrosoftRTTIProcessor::ProcessVFT(uint64_t vftAddr, ClassInfo &classInfo, std::optional<BaseClassInfo> baseClassInfo)
 {
     VirtualFunctionTableInfo vftInfo = {vftAddr};
-    // Gather all virtual functions
-    BinaryReader reader = BinaryReader(m_view);
-    reader.Seek(vftAddr);
     // Virtual functions and the analysis object of it, if it exists.
     std::vector<std::pair<uint64_t, std::optional<Ref<Function>>>> virtualFunctions = {};
+    uint64_t currentVftEntry = vftAddr;
     while (true)
     {
-        uint64_t readOffset = reader.GetOffset();
-        if (!m_view->IsValidOffset(readOffset))
+        uint64_t vFuncAddr = 0;
+        const FunctionDiscoverState state = DiscoverVirtualFunction(currentVftEntry, vFuncAddr);
+        if (state == FunctionDiscoverState::Failed)
             break;
-        uint64_t vFuncAddr = reader.ReadPointer();
-        auto funcs = m_view->GetAnalysisFunctionsForAddress(vFuncAddr);
-        if (funcs.empty())
-        {
-            Ref<Segment> segment = m_view->GetSegmentAt(vFuncAddr);
-            if (segment == nullptr || !(segment->GetFlags() & (SegmentExecutable | SegmentDenyWrite)))
-            {
-                // Last CompleteObjectLocator or hit the next CompleteObjectLocator
-                break;
-            }
-            // TODO: Is likely a function check here?
-            m_logger->LogDebugF("Discovered function from virtual function table... {:#x}", vFuncAddr);
-            auto vftPlatform = m_view->GetDefaultPlatform()->GetAssociatedPlatformByAddress(vFuncAddr);
-            auto vFunc = m_view->AddFunctionForAnalysis(vftPlatform, vFuncAddr, true);
-            virtualFunctions.emplace_back(vFuncAddr, vFunc ? std::optional(vFunc) : std::nullopt);
-        }
-        else
-        {
-            // Only ever add one function.
-            virtualFunctions.emplace_back(vFuncAddr, funcs.front());
-        }
+        currentVftEntry += m_view->GetAddressSize();
+        Ref<Platform> vftPlatform = m_view->GetDefaultPlatform()->GetAssociatedPlatformByAddress(vFuncAddr);
+        Ref<Function> vFunc = m_view->GetAnalysisFunction(vftPlatform, vFuncAddr);
+        virtualFunctions.emplace_back(vFuncAddr, vFunc ? std::optional(vFunc) : std::nullopt);
     }
 
     if (virtualFunctions.empty())
diff --git a/plugins/rtti/rtti.cpp b/plugins/rtti/rtti.cpp
index 23b13fde..fde5a6ab 100644
--- a/plugins/rtti/rtti.cpp
+++ b/plugins/rtti/rtti.cpp
@@ -4,6 +4,22 @@ using namespace BinaryNinja;
 using namespace BinaryNinja::RTTI;
 
 
+Ref<Symbol> RTTI::GetRealSymbol(BinaryView *view, uint64_t relocAddr, uint64_t symAddr)
+{
+    if (view->IsOffsetExternSemantics(symAddr))
+    {
+        // Because bases in the extern section are not 8 byte width only they will
+        // overlap with other externs, until https://github.com/Vector35/binaryninja-api/issues/6387 is fixed.
+        // Check relocation at objectAddr for symbol
+        for (const auto& r : view->GetRelocationsAt(relocAddr))
+            if (auto relocSym = r->GetSymbol())
+                return relocSym;
+    }
+
+    return view->GetSymbolByAddress(symAddr);
+}
+
+
 std::optional<std::string> RTTI::DemangleNameMS(BinaryView* view, bool allowMangled, const std::string &mangledName)
 {
     QualifiedName demangledName = {};
@@ -213,6 +229,59 @@ Ref<Metadata> RTTIProcessor::SerializedMetadata()
 }
 
 
+bool RTTIProcessor::IsLikelyFunction(uint64_t addr) const
+{
+    // Disassemble to just make a little extra certain this is a function.
+    auto vftPlatform = m_view->GetDefaultPlatform()->GetAssociatedPlatformByAddress(addr);
+    Ref<Architecture> arch = vftPlatform->GetArchitecture();
+    const size_t maxInstrLen = arch->GetMaxInstructionLength();
+    DataBuffer instrBuffer = m_view->ReadBuffer(addr, maxInstrLen);
+    InstructionInfo instrInfo;
+    const bool validInstr = arch->GetInstructionInfo(static_cast<uint8_t*>(instrBuffer.GetData()), addr, maxInstrLen, instrInfo);
+    return validInstr;
+}
+
+RTTIProcessor::FunctionDiscoverState RTTIProcessor::DiscoverVirtualFunction(uint64_t vftEntryAddr, uint64_t& vFuncAddr)
+{
+    if (!m_view->IsValidOffset(vftEntryAddr))
+        return FunctionDiscoverState::Failed;
+    BinaryReader reader = BinaryReader(m_view);
+    reader.Seek(vftEntryAddr);
+    vFuncAddr = reader.ReadPointer();
+    auto funcs = m_view->GetAnalysisFunctionsForAddress(vFuncAddr);
+    if (!funcs.empty())
+        return FunctionDiscoverState::AlreadyExists;
+
+    // Handle external virtual functions, we won't have a backing function for them.
+    if (!m_view->IsOffsetCodeSemantics(vFuncAddr))
+    {
+        // TODO: Sometimes vFunc idx will be zeroed iirc.
+        // We allow vfuncs to point to extern functions.
+        // TODO: Until https://github.com/Vector35/binaryninja-api/issues/5982 is fixed we need to check extern sym relocs instead of the symbol directly
+        auto vFuncSym = GetRealSymbol(m_view, reader.GetOffset(), vFuncAddr);
+        if (!vFuncSym)
+            return FunctionDiscoverState::Failed;
+        DataVariable dv;
+        bool foundDv = m_view->GetDataVariableAtAddress(vFuncAddr, dv);
+        // Last virtual function, or hit the next vtable.
+        if (!foundDv || !dv.type->m_object)
+            return FunctionDiscoverState::Failed;
+        // Void externs are very likely to be a func.
+        // TODO: Add some sanity checks for this!
+        if (!dv.type->IsFunction() && !(dv.type->IsVoid() && vFuncSym->GetType() == ExternalSymbol))
+            return FunctionDiscoverState::Failed;
+        return FunctionDiscoverState::Extern;
+    }
+
+    if (!IsLikelyFunction(vFuncAddr))
+        return FunctionDiscoverState::Failed;
+    m_logger->LogDebugF("Discovered function from virtual function table... {:#x}", vFuncAddr);
+    Ref<Platform> vftPlatform = m_view->GetDefaultPlatform()->GetAssociatedPlatformByAddress(vFuncAddr);
+    m_view->AddFunctionForAnalysis(vftPlatform, vFuncAddr, true);
+    return FunctionDiscoverState::Discovered;
+}
+
+
 void RTTIProcessor::DeserializedMetadata(RTTIProcessorType type, const Ref<Metadata> &metadata)
 {
     std::map<std::string, Ref<Metadata>> msvcMeta = metadata->GetKeyValueStore();
@@ -228,4 +297,4 @@ void RTTIProcessor::DeserializedMetadata(RTTIProcessorType type, const Ref<Metad
                 m_unhandledClassInfo[objectAddr] = classInfo;
         }
     }
-}
\ No newline at end of file
+}
diff --git a/plugins/rtti/rtti.h b/plugins/rtti/rtti.h
index b46a33a1..00b105eb 100644
--- a/plugins/rtti/rtti.h
+++ b/plugins/rtti/rtti.h
@@ -6,6 +6,8 @@ constexpr const char *VIEW_METADATA_RTTI = "rtti";
 constexpr int RTTI_CONFIDENCE = 100;
 
 namespace BinaryNinja::RTTI {
+	Ref<Symbol> GetRealSymbol(BinaryView *view, uint64_t relocAddr, uint64_t symAddr);
+
 	std::optional<std::string> DemangleNameMS(BinaryView* view, bool allowMangled, const std::string &mangledName);
 
 	std::optional<std::string> DemangleNameGNU3(BinaryView* view, bool allowMangled, const std::string &mangledName);
@@ -69,6 +71,14 @@ namespace BinaryNinja::RTTI {
 	class RTTIProcessor
 	{
 	protected:
+		enum class FunctionDiscoverState
+		{
+			Failed = 0,
+			AlreadyExists = 1,
+			Discovered = 2,
+			Extern = 3
+		};
+
 		Ref<BinaryView> m_view;
 		Ref<Logger> m_logger;
 
@@ -78,6 +88,10 @@ namespace BinaryNinja::RTTI {
 		virtual std::optional<ClassInfo> ProcessRTTI(uint64_t objectAddr) = 0;
 
 		virtual std::optional<VirtualFunctionTableInfo> ProcessVFT(uint64_t vftAddr, ClassInfo &classInfo, std::optional<BaseClassInfo> baseClassInfo) = 0;
+
+		[[nodiscard]] bool IsLikelyFunction(uint64_t addr) const;
+
+		[[nodiscard]] FunctionDiscoverState DiscoverVirtualFunction(uint64_t vftEntryAddr, uint64_t& vFuncAddr);
 	public:
 		virtual ~RTTIProcessor() = default;