From a4465a1f88bcd4a9626c34474ad9726bd5a30833 Mon Sep 17 00:00:00 2001
From: kat <kat@vector35.com>
Date: Mon, 10 Apr 2023 07:25:07 -0400
Subject: [PATCH] Gate ARC Cleanup behind a setting (default true)

---
 MessageHandler.cpp |  2 ++
 MessageHandler.h   |  3 +++
 Plugin.cpp         | 12 ++++++++++++
 Workflow.cpp       |  2 +-
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/MessageHandler.cpp b/MessageHandler.cpp
index 9044abe..32b97a2 100644
--- a/MessageHandler.cpp
+++ b/MessageHandler.cpp
@@ -14,6 +14,8 @@ const std::set<std::string> arcFunctionNames = {
 MessageHandler::MessageHandler(Ref<BinaryView> data)
     : m_data(data)
 {
+    m_shouldCleanupARCCode = BinaryNinja::Settings::Instance()->Get<bool>("objc.cleanupARCCode");
+
     std::unique_lock<std::recursive_mutex> lock(m_stubMutex);
 
     m_authStubsSection = data->GetSectionByName("__auth_stubs");
diff --git a/MessageHandler.h b/MessageHandler.h
index e150e93..abd790a 100644
--- a/MessageHandler.h
+++ b/MessageHandler.h
@@ -5,6 +5,7 @@
 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;
@@ -32,6 +33,8 @@ class MessageHandler : public BinaryNinja::BinaryDataNotification {
 
     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);
diff --git a/Plugin.cpp b/Plugin.cpp
index 32ae030..c6093f8 100644
--- a/Plugin.cpp
+++ b/Plugin.cpp
@@ -24,6 +24,18 @@ BINARYNINJAPLUGIN bool CorePluginInit()
     Workflow::registerActivities();
     Commands::registerCommands();
 
+    BinaryNinja::Ref<BinaryNinja::Settings> settings = BinaryNinja::Settings::Instance();
+    settings->RegisterGroup("objc", "Objective-C");
+
+    settings->RegisterSetting("objc.cleanupARCCode",
+	R"({
+	"title" : "ARC Cleanup",
+	"type" : "boolean",
+	"default" : true,
+	"description" : "Remove ARC related code, i.e. calls to _objc_release, _objc_retain, and other ARC functions, from ILs"
+	})");
+
+
     std::vector<BinaryNinja::Ref<BinaryNinja::Architecture>> targets = {
         BinaryNinja::Architecture::GetByName("aarch64"),
         BinaryNinja::Architecture::GetByName("x86_64")
diff --git a/Workflow.cpp b/Workflow.cpp
index c9b7a8f..5839706 100644
--- a/Workflow.cpp
+++ b/Workflow.cpp
@@ -219,7 +219,7 @@ void Workflow::inlineMethodCalls(AnalysisContextRef ac)
                     && params[0].operation == LLIL_REG_SSA
                     && params[1].operation == LLIL_REG_SSA)
                     rewriteMethodCall(ssa, insnIndex);
-            } else if (messageHandler->isARCFunction(callExpr.GetValue().value)) {
+            } else if (messageHandler->isARCFunction(callExpr.GetValue().value) && messageHandler->ShouldCleanupARCCode()) {
                 auto nonSSAIdx = ssa->GetNonSSAInstructionIndex(insnIndex);
                 auto targetInsn = llil->GetInstruction(nonSSAIdx);
                 if (insn.operation == LLIL_CALL_SSA)