//
// Created by kat on 8/6/24.
//

// TODO We could use an LLIL/MLIL workflow to rewrite off-image value-loads
//  	(i.e. MLIL_VAR_LOAD.MLIL_DEREF.MLIL_CONST_PTR) to just read the value out of the cache and replace the load
// 		in stub regions.
//
// This is a pretty rough workflow and has huge room for improvements all around.

#include "SharedCacheWorkflow.h"
#include "lowlevelilinstruction.h"
#include "mediumlevelilinstruction.h"
#include "../api/sharedcacheapi.h"
#include "thread"
#include <shared_mutex>

using namespace SharedCacheAPI;

struct GlobalWorkflowState
{
	std::mutex imageLoadMutex;
	bool autoLoadStubsAndDyldData = true;
	bool autoLoadObjCStubRequirements = true;
};

static std::unordered_map<uint64_t, std::shared_ptr<GlobalWorkflowState>> globalWorkflowState;
std::shared_mutex globalWorkflowStateMutex;

std::shared_ptr<GlobalWorkflowState> GetGlobalWorkflowState(Ref<BinaryView> view)
{
	std::shared_lock<std::shared_mutex> readLock(globalWorkflowStateMutex);
	uint64_t viewId = view->GetFile()->GetSessionId();
	auto foundState = globalWorkflowState.find(viewId);
	if (foundState != globalWorkflowState.end())
		return foundState->second;
	readLock.unlock();

	std::unique_lock<std::shared_mutex> writeLock(globalWorkflowStateMutex);
	globalWorkflowState[viewId] = std::make_shared<GlobalWorkflowState>();
	Ref<Settings> settings = view->GetLoadSettings(VIEW_NAME);
	bool autoLoadStubsAndDyldData = true;
	if (settings && settings->Contains("loader.dsc.autoLoadStubsAndDyldData"))
	{
		autoLoadStubsAndDyldData = settings->Get<bool>("loader.dsc.autoLoadStubsAndDyldData", view);
	}
	globalWorkflowState[viewId]->autoLoadStubsAndDyldData = autoLoadStubsAndDyldData;
	bool autoLoadObjC = true;
	if (settings && settings->Contains("loader.dsc.autoLoadObjCStubRequirements"))
	{
		autoLoadObjC = settings->Get<bool>("loader.dsc.autoLoadObjCStubRequirements", view);
	}
	globalWorkflowState[viewId]->autoLoadObjCStubRequirements = autoLoadObjC;
	return globalWorkflowState[viewId];
}


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 ProcessStubImageCall(const std::shared_ptr<GlobalWorkflowState> &workflowState, Ref<SharedCache> cache, Ref<Function> func,
						   const Ref<Section>& section, uint64_t target)
{
	if (!workflowState->imageLoadMutex.try_lock())
		return;

	auto view = func->GetView();
	if (!view->IsValidOffset(target))
	{
		if (!cache->GetImageNameForAddress(target).empty())
			cache->LoadImageContainingAddress(target);
		else
			cache->LoadSectionAtAddress(target);

		// Check to see if there are any functions inside the newly added image.
		bool newFunctions = false;
		for (const auto &sectFunc : view->GetAnalysisFunctionList())
		{
			if (section->GetStart() <= sectFunc->GetStart() && sectFunc->GetStart() < section->GetEnd())
				newFunctions = true;
		}

		// If there are any new functions we should re-analyze the current function.
		if (newFunctions)
			func->Reanalyze();
	}

	workflowState->imageLoadMutex.unlock();
}


void SharedCacheWorkflow::ProcessOffImageCall(Ref<AnalysisContext> ctx, Ref<SharedCache> cache, Ref<Function> func,
                                              Ref<MediumLevelILFunction> mssa, const MediumLevelILInstruction dest,
                                              bool applySymbolIfFoundToCurrentFunction)
{
	auto bv = func->GetView();
	WorkerPriorityEnqueue([bv=std::move(bv), cache=std::move(cache), dest=dest, func=func, applySymbolIfFoundToCurrentFunction](){
			auto workflowState = GetGlobalWorkflowState(bv);
			if (dest.operation != MLIL_CONST_PTR && dest.operation != MLIL_CONST)
				return;
			if (workflowState->autoLoadStubsAndDyldData &&
					(cache->GetNameForAddress(dest.GetConstant()).find("dyld_shared_cache_branch_islands") != std::string::npos
						|| cache->GetNameForAddress(dest.GetConstant()).find("::_stubs") != std::string::npos
					)
				)

			{
				if (cache->LoadSectionAtAddress(dest.GetConstant()))
				{
					func->Reanalyze();
				}
			}
			else
			{
				if (applySymbolIfFoundToCurrentFunction)
					cache->FindSymbolAtAddrAndApplyToAddr(dest.GetConstant(), func->GetStart(), false);
				else
					cache->FindSymbolAtAddrAndApplyToAddr(dest.GetConstant(), dest.GetConstant(), false);
			}
	});
}


void SharedCacheWorkflow::FixupStubs(Ref<AnalysisContext> ctx)
{
	try
	{
		const auto func = ctx->GetFunction();
		const auto arch = func->GetArchitecture();

		const auto bv = func->GetView();

		auto workflowState = GetGlobalWorkflowState(bv);

		auto funcStart = func->GetStart();
		auto sections = bv->GetSectionsAt(funcStart);
		if (sections.empty())
			return;
		// Just get the first section
		const auto& section = sections.front();

		const auto mlil = ctx->GetMediumLevelILFunction();
		if (!mlil)
			return;
		const auto mssa = mlil->GetSSAForm();
		if (!mssa)
			return;

		Ref<SharedCache> cache = new SharedCache(bv);
		if (cache->m_object == nullptr)
			return;

		// Processor that automatically loads the libObjC image when it encounters a stub (so we can do inlining).
		if (workflowState->autoLoadObjCStubRequirements && section->GetName().find("__objc_stubs") != std::string::npos)
		{
			auto firstInstruction = mlil->GetInstruction(0);
			if (firstInstruction.operation == MLIL_TAILCALL)
			{
				auto dest = firstInstruction.GetDestExpr<MLIL_TAILCALL>();
				if (dest.operation == MLIL_CONST_PTR)
				{
					// We're ready, everything is here
					func->SetAutoInlinedDuringAnalysis(true);
					return;
				}
			}
			for (const auto& block : mssa->GetBasicBlocks())
			{
				for (size_t i = block->GetStart(), end = block->GetEnd(); i < end; ++i)
				{
					auto instr = mssa->GetInstruction(i);
					if (instr.operation == MLIL_JUMP)
					{
						if (instr.GetDestExpr<MLIL_JUMP>().operation == MLIL_VAR_SSA)
						{
							auto dest = instr.GetDestExpr<MLIL_JUMP>();
							auto value = mssa->GetSSAVarValue(dest.GetSourceSSAVariable());
							if (value.state == UndeterminedValue)
							{
								auto def = mssa->GetSSAVarDefinition(dest.GetSourceSSAVariable());
								auto defInstr = mssa->GetInstruction(def);
								auto targetOffset = defInstr.GetSourceExpr().GetSourceExpr().GetConstant();
								ProcessStubImageCall(workflowState, cache, func, section, targetOffset);
							}
						}
						else if (instr.GetDestExpr<MLIL_JUMP>().operation == MLIL_CONST_PTR)
						{
							auto dest = instr.GetDestExpr<MLIL_JUMP>();
							auto targetOffset = dest.GetConstant();
							ProcessStubImageCall(workflowState, cache, func, section, targetOffset);
						}
					}
				}
			}

			return;
		}

		if (section->GetName().find("::_stubs") != std::string::npos // Branch Islands (iOS 16)
			|| section->GetName().find("dyld_shared_cache_branch_islands") != std::string::npos // Branch Islands (iOS 11-?)
			|| section->GetName().find("::__stubs") != std::string::npos // Stubs (non arm64e)
			|| section->GetName().find("::__auth_stubs") != std::string::npos // Stubs (arm64e)
			)
		{
			auto firstInstruction = mlil->GetInstruction(0);
			if (firstInstruction.operation == MLIL_TAILCALL)
			{
				auto dest = firstInstruction.GetDestExpr<MLIL_TAILCALL>();
				if (dest.operation == MLIL_CONST_PTR)
				{
					if (auto symbol = bv->GetSymbolByAddress(dest.GetConstant()))
					{
						auto newSymbol = new Symbol(FunctionSymbol, "j_" + symbol->GetRawName(), func->GetStart());
						bv->DefineUserSymbol(newSymbol);
					}
				}
			}
			else if (firstInstruction.operation == MLIL_JUMP)
			{
				auto dest = firstInstruction.GetDestExpr<MLIL_JUMP>();
				if (dest.operation == MLIL_CONST_PTR)
				{
					if (!bv->IsValidOffset(dest.GetConstant()))
					{
						ProcessOffImageCall(ctx, cache, func, mssa, dest, true);
					}
				}

				else if (dest.operation == MLIL_LOAD)
				{
					if (dest.GetSourceExpr().operation == MLIL_CONST_PTR)
					{
						dest = dest.GetSourceExpr();
						if (!bv->IsValidOffset(dest.GetConstant()))
						{
							ProcessOffImageCall(ctx, cache, func, mssa, dest);
						}
					}
				}
			}
			else
			{
				for (const auto& block : mssa->GetBasicBlocks())
				{
					for (size_t i = block->GetStart(), end = block->GetEnd(); i < end; ++i)
					{
						auto instr = mssa->GetInstruction(i);
						if (instr.operation == MLIL_JUMP)
						{
							if (instr.GetDestExpr<MLIL_JUMP>().operation == MLIL_VAR_SSA)
							{
								auto dest = instr.GetDestExpr<MLIL_JUMP>();
								auto value = mssa->GetSSAVarValue(dest.GetSourceSSAVariable());
								if (value.state == UndeterminedValue)
								{
									auto def = mssa->GetSSAVarDefinition(dest.GetSourceSSAVariable());
									auto defInstr = mssa->GetInstruction(def);
									auto targetOffset = defInstr.GetSourceExpr().GetSourceExpr().GetConstant();
									ProcessStubImageCall(workflowState, cache, func, section, targetOffset);
								}
							}
						}
					}
				}
			}

			return;
		}

		for (const auto& block : mssa->GetBasicBlocks())
		{
			for (size_t i = block->GetStart(), end = block->GetEnd(); i < end; ++i)
			{
				auto instr = mssa->GetInstruction(i);
				if (instr.operation == MLIL_CALL_SSA)
				{
					if (instr.GetDestExpr<MLIL_CALL_SSA>().operation == MLIL_CONST_PTR)
					{
						auto dest = instr.GetDestExpr<MLIL_CALL_SSA>();
						if (!bv->IsValidOffset(dest.GetConstant()))
						{
							ProcessOffImageCall(ctx, cache, func, mssa, dest);
						}
					}
				}
			}
		}
	}
	catch (...)
	{}
}


static constexpr auto workflowInfo = R"({
  "title": "Shared Cache Workflow",
  "description": "Shared Cache Workflow",
  "capabilities": []
})";


void fixObjCCallTypes(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)
		{
			// Filter out calls that aren't to `objc_msgSend`.
			auto callExpr = insn.GetDestExpr<LLIL_CALL_SSA>();
			bool isMessageSend = false;
			if (auto symbol = bv->GetSymbolByAddress(callExpr.GetValue().value))
				isMessageSend = symbol->GetRawName() == "_objc_msgSend";
			if (!isMessageSend)
				return;

			const auto llil = ssa->GetNonSSAForm();
			const auto insn = ssa->GetInstruction(insnIndex);
			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 = BinaryNinja::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 = BinaryNinja::Type::PointerType(ssa->GetArchitecture(), BinaryNinja::Type::VoidType());

			std::vector<BinaryNinja::FunctionParameter> callTypeParams;
			auto cc = bv->GetDefaultPlatform()->GetDefaultCallingConvention();

			callTypeParams.push_back({"self", retType, true, BinaryNinja::Variable()});

			auto selType = bv->GetTypeByName({ "SEL" });
			if (!selType)
				selType = BinaryNinja::Type::PointerType(ssa->GetArchitecture(), BinaryNinja::Type::IntegerType(1, true));
			callTypeParams.push_back({"sel", selType, true, BinaryNinja::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 = BinaryNinja::Type::IntegerType(bv->GetAddressSize(), true);
				if (argumentNames.size() > i && !argumentNames[i].empty())
					callTypeParams.push_back({argumentNames[i], argType, true, BinaryNinja::Variable()});
				else
					callTypeParams.push_back({"arg" + std::to_string(i), argType, true, BinaryNinja::Variable()});
			}

			auto funcType = BinaryNinja::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);
}



void SharedCacheWorkflow::Register()
{
	Ref<Workflow> wf = BinaryNinja::Workflow::Instance("core.function.baseAnalysis")->Clone("core.function.dsc");
	wf->RegisterActivity(new BinaryNinja::Activity("core.analysis.dscstubs", &SharedCacheWorkflow::FixupStubs));
	wf->RegisterActivity(new BinaryNinja::Activity("core.analysis.fixObjCCallTypes", &fixObjCCallTypes));
	wf->Insert("core.function.analyzeTailCalls", "core.analysis.fixObjCCallTypes");
	wf->Insert("core.function.analyzeTailCalls", "core.analysis.dscstubs");

	BinaryNinja::Workflow::RegisterWorkflow(wf, workflowInfo);
}

extern "C"
{
	void RegisterSharedCacheWorkflow()
	{
		SharedCacheWorkflow::Register();
	}
}