Remove ARC-related Obj-C Code from decompilation

MessageHandler.cpp

@@ -2,17 +2,90 @@
 
 using namespace BinaryNinja;
 
+const std::set<std::string> arcFunctionNames = {
+    "_objc_retain",
+    "_objc_release",
+    "_objc_autorelease",
+    "_objc_autoreleaseReturnValue",
+    "_objc_retainAutoreleasedReturnValue",
+    "_objc_unsafeClaimAutoreleasedReturnValue"
+};
+
 MessageHandler::MessageHandler(Ref<BinaryView> data)
+    : m_data(data)
 {
+    m_shouldCleanupARCCode = BinaryNinja::Settings::Instance()->Get<bool>("workflows.objectiveC.cleanupARCCode");
+
+    std::unique_lock<std::recursive_mutex> lock(m_stubMutex);
+
+    m_authStubsSection = data->GetSectionByName("__auth_stubs");
+    m_stubsSection = data->GetSectionByName("__stubs");
+
+    if (!m_authStubsSection && !m_stubsSection)
+        m_readyForRealAnalysisPass = true;
+    else
+        data->RegisterNotification(this);
+
     m_msgSendFunctions = findMsgSendFunctions(data);
+    m_arcFunctions = findARCFunctions(data);
+}
+
+void MessageHandler::OnSymbolAdded(BinaryNinja::BinaryView* view, BinaryNinja::Symbol* sym)
+{
+    std::unique_lock<std::recursive_mutex> lock(m_stubMutex);
+
+    if (m_readyForRealAnalysisPass)
+        return;
+    if (sym->GetType() != ImportedFunctionSymbol)
+        return;
+    if (!m_presentRequiredStubs.count(sym->GetShortName()))
+        return;
+
+    m_locatedRequiredStubs.insert(sym->GetShortName());
+    if (arcFunctionNames.count(sym->GetShortName()))
+        m_arcFunctions.insert(sym->GetAddress());
+    else if (sym->GetShortName() == "_objc_msgSend") {
+        auto f = m_data->GetAnalysisFunction(m_data->GetDefaultPlatform(), sym->GetAddress());
+        if (!f)
+            return;
+        auto retType = BinaryNinja::Confidence<BinaryNinja::Ref<BinaryNinja::Type>>(
+            BinaryNinja::Type::PointerType(m_data->GetAddressSize(), BinaryNinja::Type::VoidType(),
+                0));
+
+        std::vector<BinaryNinja::FunctionParameter> params;
+        auto cc = m_data->GetDefaultPlatform()->GetDefaultCallingConvention();
+
+        params.push_back({ "self",
+            BinaryNinja::Type::NamedType(m_data, { "id" }),
+            true,
+            BinaryNinja::Variable() });
+        params.push_back({ "sel",
+            BinaryNinja::Type::PointerType(m_data->GetAddressSize(), BinaryNinja::Type::IntegerType(1, false)),
+            true,
+            BinaryNinja::Variable() });
+
+        auto funcType = BinaryNinja::Type::FunctionType(retType, cc, params, true);
+        f->SetUserType(funcType);
+
+        m_msgSendFunctions.insert(f->GetStart());
+    }
+
+    if (m_locatedRequiredStubs.size() == m_presentRequiredStubs.size()) {
+        m_readyForRealAnalysisPass = true;
+
+        for (auto fAddr : m_reAnalysisRequiredFunctions)
+            for (auto f : m_data->GetAnalysisFunctionsForAddress(fAddr))
+                f->Reanalyze();
+
+        m_data->UnregisterNotification(this);
+    }
 }
 
 std::set<uint64_t> MessageHandler::findMsgSendFunctions(BinaryNinja::Ref<BinaryNinja::BinaryView> data)
 {
-    std::set<uint64_t> results;
+    std::unique_lock<std::recursive_mutex> lock(m_stubMutex);
 
-    const auto authStubsSection = data->GetSectionByName("__auth_stubs");
-    const auto stubsSection = data->GetSectionByName("__stubs");
+    std::set<uint64_t> results;
     const auto authGotSection = data->GetSectionByName("__auth_got");
     const auto gotSection = data->GetSectionByName("__got");
     const auto laSymbolPtrSection = data->GetSectionByName("__la_symbol_ptr");
@@ -26,38 +99,95 @@ std::set<uint64_t> MessageHandler::findMsgSendFunctions(BinaryNinja::Ref<BinaryN
         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))
+        if ((authGotSection && sectionContains(authGotSection, c))
+            || (gotSection && sectionContains(gotSection, c))
+            || (laSymbolPtrSection && sectionContains(laSymbolPtrSection, c))) {
+            results.insert(c->GetAddress());
+        }
+        if ((m_authStubsSection && sectionContains(m_authStubsSection, c))
+            || (m_stubsSection && sectionContains(m_stubsSection, c))) {
+            results.insert(c->GetAddress());
+            m_locatedRequiredStubs.insert(c->GetShortName());
+        }
+    }
+    m_presentRequiredStubs.insert("_objc_msgSend");
+
+    return results;
+}
+
+std::set<uint64_t> MessageHandler::findARCFunctions(BinaryNinja::Ref<BinaryNinja::BinaryView> data)
+{
+    std::unique_lock<std::recursive_mutex> lock(m_stubMutex);
+
+    std::set<uint64_t> results;
+    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;
+    };
+
+    std::vector<BinaryNinja::Ref<BinaryNinja::Symbol>> candidates;
+    std::vector<BinaryNinja::Ref<BinaryNinja::Symbol>> next;
+    for (auto name : arcFunctionNames) {
+        next = data->GetSymbolsByName(name);
+        candidates.insert(candidates.end(), next.begin(), next.end());
+    }
+    for (auto& c : candidates) {
+        if ((authGotSection && sectionContains(authGotSection, c))
             || (gotSection && sectionContains(gotSection, c))
             || (laSymbolPtrSection && sectionContains(laSymbolPtrSection, c))) {
+            {
+                results.insert(c->GetAddress());
+                m_presentRequiredStubs.insert(c->GetShortName());
+            }
+        }
+        if ((m_authStubsSection && sectionContains(m_authStubsSection, c))
+            || (m_stubsSection && sectionContains(m_stubsSection, c))) {
+            m_locatedRequiredStubs.insert(c->GetShortName());
             results.insert(c->GetAddress());
         }
     }
 
     return results;
 }
 
+void MessageHandler::functionWasAnalyzed(uint64_t addr)
+{
+    if (!m_readyForRealAnalysisPass) {
+        std::unique_lock<std::mutex> lock(m_reAnalysisRequiredFunctionsMutex);
+        m_reAnalysisRequiredFunctions.insert(addr);
+    }
+}
+
 bool MessageHandler::isMessageSend(uint64_t functionAddress)
 {
+    std::unique_lock<std::recursive_mutex> lock(m_stubMutex);
     return m_msgSendFunctions.count(functionAddress);
 }
+
+bool MessageHandler::isARCFunction(uint64_t functionAddress)
+{
+    std::unique_lock<std::recursive_mutex> lock(m_stubMutex);
+    return m_arcFunctions.count(functionAddress);
+}
+
+bool MessageHandler::isFunctionLocatedInStubSection(uint64_t addr)
+{
+    if (m_stubsSection && m_stubsSection->GetStart() < addr
+        && m_stubsSection->GetStart() + m_stubsSection->GetLength() > addr)
+        return true;
+    else if (m_authStubsSection && m_authStubsSection->GetStart() < addr
+        && m_authStubsSection->GetStart() + m_authStubsSection->GetLength() > addr)
+        return true;
+
+    return false;
+}

MessageHandler.h

@@ -2,15 +2,43 @@
 
 #include <binaryninjaapi.h>
 
-class MessageHandler {
+class MessageHandler : public BinaryNinja::BinaryDataNotification {
+
+    BinaryNinja::Ref<BinaryNinja::BinaryView> m_data;
+    bool m_shouldCleanupARCCode;
+
+    BinaryNinja::Ref<BinaryNinja::Section> m_authStubsSection;
+    BinaryNinja::Ref<BinaryNinja::Section> m_stubsSection;
+
+    bool m_readyForRealAnalysisPass = false;
+    std::recursive_mutex m_stubMutex;
+    std::set<std::string> m_presentRequiredStubs;
+    std::set<std::string> m_locatedRequiredStubs;
+
+    std::mutex m_reAnalysisRequiredFunctionsMutex;
+    std::set<uint64_t> m_reAnalysisRequiredFunctions;
 
     std::set<uint64_t> m_msgSendFunctions;
-    static std::set<uint64_t> findMsgSendFunctions(BinaryNinja::Ref<BinaryNinja::BinaryView> data);
+    std::set<uint64_t> m_checkedNonMsgSendFunctions;
+    std::set<uint64_t> m_arcFunctions;
+    std::set<uint64_t> m_checkedNonARCFunctions;
+
+    std::set<uint64_t> findMsgSendFunctions(BinaryNinja::Ref<BinaryNinja::BinaryView> data);
+    std::set<uint64_t> findARCFunctions(BinaryNinja::Ref<BinaryNinja::BinaryView> data);
+
+    virtual void OnSymbolAdded(BinaryNinja::BinaryView* view, BinaryNinja::Symbol* sym) override;
 
 public:
     MessageHandler(BinaryNinja::Ref<BinaryNinja::BinaryView> data);
 
+    void functionWasAnalyzed(uint64_t addr);
+
+    bool ShouldCleanupARCCode() const { return m_shouldCleanupARCCode; }
+
     std::set<uint64_t> getMessageSendFunctions() const { return m_msgSendFunctions; }
     bool hasMessageSendFunctions() const { return m_msgSendFunctions.size() != 0; }
     bool isMessageSend(uint64_t);
+    bool isARCFunction(uint64_t);
+
+    bool isFunctionLocatedInStubSection(uint64_t);
 };

Plugin.cpp

@@ -31,6 +31,18 @@ BINARYNINJAPLUGIN bool CorePluginInit()
     Workflow::registerActivities();
     Commands::registerCommands();
 
+    BinaryNinja::Ref<BinaryNinja::Settings> settings = BinaryNinja::Settings::Instance();
+    settings->RegisterGroup("objc", "Objective-C");
+
+    settings->RegisterSetting("workflows.objectiveC.cleanupARCCode",
+	R"({
+	"title" : "Hide ARC Calls",
+	"type" : "boolean",
+	"default" : true,
+	"description" : "Remove ARC-related code, e.g. calls to _objc_release, _objc_retain, and other ARC functions, in ILs"
+	})");
+
+
     std::vector<BinaryNinja::Ref<BinaryNinja::Architecture>> targets = {
         BinaryNinja::Architecture::GetByName("aarch64"),
         BinaryNinja::Architecture::GetByName("x86_64")