1 files changed, 129 insertions, 0 deletions
diff --git a/view/sharedcache/workflow/ObjCActivity.cpp b/view/sharedcache/workflow/ObjCActivity.cpp
new file mode 100644
index 00000000..55a8a8fa
--- /dev/null
+++ b/view/sharedcache/workflow/ObjCActivity.cpp
@@ -0,0 +1,129 @@
+#include "ObjCActivity.h"
+#include "lowlevelilinstruction.h"
+
+// TODO: Consolidate this with the Obj-C workflow at some point https://github.com/Vector35/workflow_objc
+
+using namespace BinaryNinja;
+
+void ObjCActivity::Register(Workflow &workflow)
+{
+    workflow.RegisterActivity(new Activity("core.analysis.objc.adjustCallType", &AdjustCallType));
+	workflow.Insert("core.function.analyzeTailCalls", "core.analysis.objc.adjustCallType");
+}
+
+std::vector<std::string> splitSelector(const std::string& selector) {
+	std::vector<std::string> components;
+	std::istringstream stream(selector);
+	std::string component;
+
+	while (std::getline(stream, component, ':')) {
+		if (!component.empty()) {
+			components.push_back(component);
+		}
+	}
+
+	return components;
+}
+
+std::vector<std::string> generateArgumentNames(const std::vector<std::string>& components) {
+	std::vector<std::string> argumentNames;
+
+	for (const std::string& component : components) {
+		size_t startPos = component.find_last_of(' ');
+		std::string argumentName = (startPos == std::string::npos) ? component : component.substr(startPos + 1);
+		argumentNames.push_back(argumentName);
+	}
+
+	return argumentNames;
+}
+
+void ObjCActivity::AdjustCallType(Ref<AnalysisContext> ctx)
+{
+    const auto func = ctx->GetFunction();
+	const auto arch = func->GetArchitecture();
+	const auto bv = func->GetView();
+
+	const auto llil = ctx->GetLowLevelILFunction();
+	if (!llil) {
+		return;
+	}
+	const auto ssa = llil->GetSSAForm();
+	if (!ssa) {
+		return;
+	}
+
+	const auto rewriteIfEligible = [bv, ssa](size_t insnIndex) {
+		auto insn = ssa->GetInstruction(insnIndex);
+		if (insn.operation != LLIL_CALL_SSA)
+			return;
+
+		// Filter out calls that aren't to `objc_msgSend`.
+		auto callExpr = insn.GetDestExpr<LLIL_CALL_SSA>();
+		if (auto symbol = bv->GetSymbolByAddress(callExpr.GetValue().value))
+			if (symbol->GetRawName() != "_objc_msgSend")
+				return;
+
+		const auto params = insn.GetParameterExprs<LLIL_CALL_SSA>();
+		// The second parameter passed to the objc_msgSend call is the address of
+		// either the selector reference or the method's name, which in both cases
+		// is dereferenced to retrieve a selector.
+		if (params.size() < 2)
+			return;
+		uint64_t rawSelector = 0;
+		if (params[1].operation == LLIL_REG_SSA)
+		{
+			const auto selectorRegister = params[1].GetSourceSSARegister<LLIL_REG_SSA>();
+			rawSelector = ssa->GetSSARegisterValue(selectorRegister).value;
+		}
+		else if (params[0].operation == LLIL_SEPARATE_PARAM_LIST_SSA)
+		{
+			if (params[0].GetParameterExprs<LLIL_SEPARATE_PARAM_LIST_SSA>().size() == 0)
+				return;
+			const auto selectorRegister = params[0].GetParameterExprs<LLIL_SEPARATE_PARAM_LIST_SSA>()[1].GetSourceSSARegister<LLIL_REG_SSA>();
+			rawSelector = ssa->GetSSARegisterValue(selectorRegister).value;
+		}
+		if (!rawSelector || !bv->IsValidOffset(rawSelector))
+			return;
+
+		// -- Do callsite override
+		auto reader = BinaryReader(bv);
+		reader.Seek(rawSelector);
+		auto selector = reader.ReadCString(500);
+		auto additionalArgumentCount = std::count(selector.begin(), selector.end(), ':');
+
+		auto retType = bv->GetTypeByName({ "id" });
+		if (!retType)
+			retType = Type::PointerType(ssa->GetArchitecture(), Type::VoidType());
+
+		std::vector<FunctionParameter> callTypeParams;
+		auto cc = bv->GetDefaultPlatform()->GetDefaultCallingConvention();
+
+		callTypeParams.emplace_back("self", retType, true, Variable());
+
+		auto selType = bv->GetTypeByName({ "SEL" });
+		if (!selType)
+			selType = Type::PointerType(ssa->GetArchitecture(), Type::IntegerType(1, true));
+		callTypeParams.emplace_back("sel", selType, true, Variable());
+
+		std::vector<std::string> selectorComponents = splitSelector(selector);
+		std::vector<std::string> argumentNames = generateArgumentNames(selectorComponents);
+
+		for (size_t i = 0; i < additionalArgumentCount; i++)
+		{
+			auto argType = Type::IntegerType(bv->GetAddressSize(), true);
+			if (argumentNames.size() > i && !argumentNames[i].empty())
+				callTypeParams.emplace_back(argumentNames[i], argType, true, Variable());
+			else
+				callTypeParams.emplace_back("arg" + std::to_string(i), argType, true, Variable());
+		}
+
+		auto funcType = Type::FunctionType(retType, cc, callTypeParams);
+		ssa->GetFunction()->SetAutoCallTypeAdjustment(ssa->GetFunction()->GetArchitecture(), insn.address, {funcType, BN_DEFAULT_CONFIDENCE});
+		// --
+	};
+
+	for (const auto& block : ssa->GetBasicBlocks())
+		for (size_t i = block->GetStart(), end = block->GetEnd(); i < end; ++i)
+			rewriteIfEligible(i);
+}
+