Move msgSend func lookups to a cache class

Author: 0cyn

CMakeLists.txt

@@ -46,6 +46,8 @@ set(PLUGIN_SOURCE
   GlobalState.cpp
   InfoHandler.h
   InfoHandler.cpp
+  MessageHandler.cpp
+  MessageHandler.h
   Plugin.cpp
   Workflow.h
   Workflow.cpp)

GlobalState.cpp

@@ -11,8 +11,20 @@
 #include <unordered_map>
 
 static std::unordered_map<BinaryViewID, SharedAnalysisInfo> g_analysisRecords;
+static std::unordered_map<BinaryViewID, MessageHandler*> g_messageHandlers;
 static std::set<BinaryViewID> g_ignoredViews;
 
+MessageHandler* GlobalState::messageHandler(BinaryViewRef bv)
+{
+    if (auto messageHandler = g_messageHandlers.find(id(bv)); messageHandler != g_messageHandlers.end()) {
+        return messageHandler->second;
+    } else {
+        auto newMessageHandler = new MessageHandler(bv);
+        g_messageHandlers[id(bv)] = newMessageHandler;
+        return newMessageHandler;
+    }
+}
+
 BinaryViewID GlobalState::id(BinaryViewRef bv)
 {
     return bv->GetFile()->GetSessionId();

GlobalState.h

@@ -10,6 +10,7 @@
 #include "BinaryNinja.h"
 
 #include "Core/AnalysisInfo.h"
+#include "MessageHandler.h"
 
 using SharedAnalysisInfo = std::shared_ptr<ObjectiveNinja::AnalysisInfo>;
 
@@ -38,6 +39,10 @@ class GlobalState {
      */
     static SharedAnalysisInfo analysisInfo(BinaryViewRef);
 
+    /**
+     * Get ObjC Message Handler for a view
+     */
+    static MessageHandler* messageHandler(BinaryViewRef);
     /**
      * Store analysis info for a view.
      */

MessageHandler.cpp

@@ -0,0 +1,63 @@
+#include "MessageHandler.h"
+
+using namespace BinaryNinja;
+
+MessageHandler::MessageHandler(Ref<BinaryView> data)
+{
+    m_msgSendFunctions = findMsgSendFunctions(data);
+}
+
+std::set<uint64_t> MessageHandler::findMsgSendFunctions(BinaryNinja::Ref<BinaryNinja::BinaryView> data)
+{
+    std::set<uint64_t> results;
+
+    const auto authStubsSection = data->GetSectionByName("__auth_stubs");
+    const auto stubsSection = data->GetSectionByName("__stubs");
+    const auto authGotSection = data->GetSectionByName("__auth_got");
+    const auto gotSection = data->GetSectionByName("__got");
+    const auto laSymbolPtrSection = data->GetSectionByName("__la_symbol_ptr");
+
+    // Shorthand to check if a symbol lies in a given section.
+    auto sectionContains = [](Ref<Section> section, Ref<Symbol> symbol) {
+        const auto start = section->GetStart();
+        const auto length = section->GetLength();
+        const auto address = symbol->GetAddress();
+
+        return (uint64_t)(address - start) <= length;
+    };
+
+    // There can be multiple `_objc_msgSend` symbols in the same binary; there
+    // may even be lots. Some of them are valid, others aren't. In order of
+    // preference, `_objc_msgSend` symbols in the following sections are
+    // preferred:
+    //
+    //   1. __auth_stubs
+    //   2. __stubs
+    //   3. __auth_got
+    //   4. __got
+    //   ?. __la_symbol_ptr
+    //
+    // There is often an `_objc_msgSend` symbol that is a stub function, found
+    // in the `__stubs` section, which will come with an imported symbol of the
+    // same name in the `__got` section. Not all `__objc_msgSend` calls will be
+    // routed through the stub function, making it important to make note of
+    // both symbols' addresses. Furthermore, on ARM64, the `__auth{stubs,got}`
+    // sections are preferred over their unauthenticated counterparts.
+    const auto candidates = data->GetSymbolsByName("_objc_msgSend");
+    for (const auto& c : candidates) {
+        if ((authStubsSection && sectionContains(authStubsSection, c))
+            || (stubsSection && sectionContains(stubsSection, c))
+            || (authGotSection && sectionContains(authGotSection, c))
+            || (gotSection && sectionContains(gotSection, c))
+            || (laSymbolPtrSection && sectionContains(laSymbolPtrSection, c))) {
+            results.insert(c->GetAddress());
+        }
+    }
+
+    return results;
+}
+
+bool MessageHandler::isMessageSend(uint64_t functionAddress)
+{
+    return m_msgSendFunctions.count(functionAddress);
+}

MessageHandler.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <binaryninjaapi.h>
+
+class MessageHandler {
+
+    std::set<uint64_t> m_msgSendFunctions;
+    static std::set<uint64_t> findMsgSendFunctions(BinaryNinja::Ref<BinaryNinja::BinaryView> data);
+
+public:
+    MessageHandler(BinaryNinja::Ref<BinaryNinja::BinaryView> data);
+
+    std::set<uint64_t> getMessageSendFunctions() const { return m_msgSendFunctions; }
+    bool hasMessageSendFunctions() const { return m_msgSendFunctions.size() != 0; }
+    bool isMessageSend(uint64_t);
+};

Workflow.cpp

@@ -26,55 +26,6 @@ static std::mutex g_initialAnalysisMutex;
 using SectionRef = BinaryNinja::Ref<BinaryNinja::Section>;
 using SymbolRef = BinaryNinja::Ref<BinaryNinja::Symbol>;
 
-std::set<uint64_t> Workflow::findMsgSendFunctions(BinaryViewRef bv)
-{
-    std::set<uint64_t> results;
-
-    const auto authStubsSection = bv->GetSectionByName("__auth_stubs");
-    const auto stubsSection = bv->GetSectionByName("__stubs");
-    const auto authGotSection = bv->GetSectionByName("__auth_got");
-    const auto gotSection = bv->GetSectionByName("__got");
-    const auto laSymbolPtrSection = bv->GetSectionByName("__la_symbol_ptr");
-
-    // Shorthand to check if a symbol lies in a given section.
-    auto sectionContains = [](SectionRef section, SymbolRef symbol) {
-        const auto start = section->GetStart();
-        const auto length = section->GetLength();
-        const auto address = symbol->GetAddress();
-
-        return (uint64_t)(address - start) <= length;
-    };
-
-    // There can be multiple `_objc_msgSend` symbols in the same binary; there
-    // may even be lots. Some of them are valid, others aren't. In order of
-    // preference, `_objc_msgSend` symbols in the following sections are
-    // preferred:
-    //
-    //   1. __auth_stubs
-    //   2. __stubs
-    //   3. __auth_got
-    //   4. __got
-    //   ?. __la_symbol_ptr
-    //
-    // There is often an `_objc_msgSend` symbol that is a stub function, found
-    // in the `__stubs` section, which will come with an imported symbol of the
-    // same name in the `__got` section. Not all `__objc_msgSend` calls will be
-    // routed through the stub function, making it important to make note of
-    // both symbols' addresses. Furthermore, on ARM64, the `__auth{stubs,got}`
-    // sections are preferred over their unauthenticated counterparts.
-    const auto candidates = bv->GetSymbolsByName("_objc_msgSend");
-    for (const auto& c : candidates) {
-        if ((authStubsSection && sectionContains(authStubsSection, c))
-            || (stubsSection && sectionContains(stubsSection, c))
-            || (authGotSection && sectionContains(authGotSection, c))
-            || (gotSection && sectionContains(gotSection, c))
-            || (laSymbolPtrSection && sectionContains(laSymbolPtrSection, c)))
-            results.insert(c->GetAddress());
-    }
-
-    return results;
-}
-
 void Workflow::rewriteMethodCall(LLILFunctionRef ssa, size_t insnIndex)
 {
     const auto bv = ssa->GetFunction()->GetView();
@@ -187,6 +138,7 @@ void Workflow::inlineMethodCalls(AnalysisContextRef ac)
             CustomTypes::defineAll(bv);
             CFStringArchitectureHook* currentHook = new CFStringArchitectureHook(bv->GetDefaultArchitecture());
             bv->GetDefaultArchitecture()->Register(currentHook);
+            auto messageHandler = GlobalState::messageHandler(bv);
 
             try {
                 auto file = std::make_shared<ObjectiveNinja::BinaryViewFile>(bv);
@@ -200,7 +152,7 @@ void Workflow::inlineMethodCalls(AnalysisContextRef ac)
 
                 InfoHandler::applyInfoToView(info, bv);
 
-                const auto msgSendFunctions = findMsgSendFunctions(bv);
+                const auto msgSendFunctions = messageHandler->getMessageSendFunctions();
                 for (auto addr : msgSendFunctions)
                 {
                     BinaryNinja::QualifiedNameAndType nameAndType;
@@ -237,12 +189,8 @@ void Workflow::inlineMethodCalls(AnalysisContextRef ac)
         }
     }
 
-    // Try to find the `objc_msgSend` functions(s), abort activity if missing.
-    //
-    // TODO: These results should be cached somehow as it can't be efficient to
-    // repeatedly search for all the usable function addresses.
-    const auto msgSendFunctions = findMsgSendFunctions(bv);
-    if (msgSendFunctions.empty()) {
+    auto messageHandler = GlobalState::messageHandler(bv);
+    if (!messageHandler->hasMessageSendFunctions()) {
         log->LogError("Cannot perform Objective-C IL cleanup; no objc_msgSend candidates found");
         GlobalState::addIgnoredView(bv);
         return;
@@ -259,14 +207,14 @@ void Workflow::inlineMethodCalls(AnalysisContextRef ac)
         return;
     }
 
-    const auto rewriteIfEligible = [bv, msgSendFunctions, ssa](size_t insnIndex) {
+    const auto rewriteIfEligible = [bv, messageHandler, ssa](size_t insnIndex) {
         auto insn = ssa->GetInstruction(insnIndex);
 
         if (insn.operation == LLIL_CALL_SSA)
         {
             // Filter out calls that aren't to `objc_msgSend`.
             auto callExpr = insn.GetDestExpr<LLIL_CALL_SSA>();
-            if (!msgSendFunctions.count(callExpr.GetValue().value))
+            if (!messageHandler->isMessageSend(callExpr.GetValue().value))
                 return;
 
             // By convention, the selector is the second argument to `objc_msgSend`,
@@ -279,6 +227,7 @@ void Workflow::inlineMethodCalls(AnalysisContextRef ac)
                 return;
 
             rewriteMethodCall(ssa, insnIndex);
+
         }
         else if (insn.operation == LLIL_SET_REG_SSA)
         {

Workflow.h

@@ -23,15 +23,6 @@ constexpr auto ResolveMethodCalls = "core.function.objectiveC.resolveMethodCalls
  */
 class Workflow {
 
-    /**
-     * Get the addresses of all usable `_objc_msgSend` functions.
-     *
-     * Not all `_objc_msgSend` functions in the binary are usable, and there
-     * may also be more than one usable candidate. For additional details and
-     * specifics, see this function's implementation.
-     */
-    static std::set<uint64_t> findMsgSendFunctions(BinaryViewRef);
-
     /**
      * Attempt to rewrite the `objc_msgSend` call at `insnIndex` with a direct
      * call to the requested method's implementation.