abb: minor refactor

Author: rssor

architecture.cpp

@@ -204,6 +204,235 @@ vector<InstructionTextToken> InstructionTextToken::ConvertInstructionTextTokenLi
 	return result;
 }
 
+BasicBlockAnalysisContext::BasicBlockAnalysisContext(BNBasicBlockAnalysisContext* context)
+{
+	m_context = context;
+}
+
+const std::map<ArchAndAddr, std::set<ArchAndAddr>> BasicBlockAnalysisContext::GetIndirectBranches()
+{
+	if (!m_indirectBranches)
+	{
+		auto& indirectBranches = m_indirectBranches.emplace();
+
+		for (size_t i = 0; i < m_context->indirectBranchesCount; i++)
+		{
+			ArchAndAddr src(new CoreArchitecture(m_context->indirectBranches[i].sourceArch), m_context->indirectBranches[i].sourceAddr);
+			ArchAndAddr dst(new CoreArchitecture(m_context->indirectBranches[i].destArch), m_context->indirectBranches[i].destAddr);
+
+			indirectBranches[src].insert(dst);
+		}
+	}
+
+	return *m_indirectBranches;
+}
+
+
+const std::set<ArchAndAddr>& BasicBlockAnalysisContext::GetIndirectNoReturnCalls()
+{
+	if (!m_indirectNoReturnCalls)
+	{
+		auto& indirectNoReturnCalls = m_indirectNoReturnCalls.emplace();
+
+		for (size_t i = 0; i < m_context->indirectNoReturnCallsCount; i++)
+		{
+			ArchAndAddr addr(new CoreArchitecture(m_context->indirectNoReturnCalls[i].arch), m_context->indirectNoReturnCalls[i].address);
+			indirectNoReturnCalls.insert(addr);
+		}
+	}
+
+	return *m_indirectNoReturnCalls;
+}
+
+
+std::map<ArchAndAddr, bool>& BasicBlockAnalysisContext::GetContextualReturns()
+{
+	if (!m_contextualReturns)
+	{
+		auto& contextualReturns = m_contextualReturns.emplace();
+
+		for (size_t i = 0; i < m_context->contextualFunctionReturnCount; i++)
+		{
+			ArchAndAddr addr(new CoreArchitecture(m_context->contextualFunctionReturnLocations[i].arch), m_context->contextualFunctionReturnLocations[i].address);
+			contextualReturns[addr] = m_context->contextualFunctionReturnValues[i];
+		}
+	}
+
+	return *m_contextualReturns;
+}
+
+
+std::map<uint64_t, std::set<ArchAndAddr>>& BasicBlockAnalysisContext::GetDirectCodeReferences()
+{
+	if (!m_directCodeReferences)
+	{
+		auto& directCodeReferences = m_directCodeReferences.emplace();
+
+		for (size_t i = 0; i < m_context->directRefCount; i++)
+		{
+			ArchAndAddr src(new CoreArchitecture(m_context->directRefSources[i].arch), m_context->directRefSources[i].address);
+
+			directCodeReferences[m_context->directRefTargets[i]].insert(src);
+		}
+	}
+
+	return *m_directCodeReferences;
+}
+
+
+std::set<ArchAndAddr>& BasicBlockAnalysisContext::GetDirectNoReturnCalls()
+{
+	if (!m_directNoReturnCalls)
+		m_directNoReturnCalls.emplace();
+
+	return *m_directNoReturnCalls;
+}
+
+
+std::set<ArchAndAddr>& BasicBlockAnalysisContext::GetHaltedDisassemblyAddresses()
+{
+	if (!m_haltedDisassemblyAddresses)
+		m_haltedDisassemblyAddresses.emplace();
+
+	return *m_haltedDisassemblyAddresses;
+}
+
+
+void BasicBlockAnalysisContext::AddTempOutgoingReference(Function* targetFunc)
+{
+	BNAnalyzeBasicBlocksContextAddTempReference(m_context, targetFunc->m_object);
+}
+
+
+Ref<BasicBlock> BasicBlockAnalysisContext::CreateBasicBlock(Architecture* arch, uint64_t start)
+{
+	BNBasicBlock* block = BNAnalyzeBasicBlocksContextCreateBasicBlock(m_context, arch->GetObject(), start);
+	if (!block)
+		return nullptr;
+	return new BasicBlock(block);
+}
+
+
+void BasicBlockAnalysisContext::AddFunctionBasicBlock(BasicBlock* block)
+{
+	BNAnalyzeBasicBlocksContextAddBasicBlockToFunction(m_context, block->GetObject());
+}
+
+
+void BasicBlockAnalysisContext::Finalize()
+{
+	if (m_directCodeReferences)
+	{
+		auto& directRefs = *m_directCodeReferences;
+
+		size_t total = 0;
+		for (auto& pair : directRefs)
+			total += pair.second.size();
+
+		BNArchitectureAndAddress* sources = new BNArchitectureAndAddress[total];
+		uint64_t* targets = new uint64_t[total];
+
+		size_t i = 0;
+		for (auto& pair : directRefs)
+		{
+			for (auto& src : pair.second)
+			{
+				sources[i].arch = src.arch->GetObject();
+				sources[i].address = src.address;
+				targets[i] = pair.first;
+				i++;
+			}
+		}
+
+		BNAnalyzeBasicBlocksContextSetDirectCodeReferences(m_context, sources, targets, total);
+
+		delete[] sources;
+		delete[] targets;
+	}
+
+	if (m_directNoReturnCalls)
+	{
+		auto& directNoReturnCalls = *m_directNoReturnCalls;
+
+		BNArchitectureAndAddress* noRets = new BNArchitectureAndAddress[directNoReturnCalls.size()];
+
+		size_t i = 0;
+		for (auto& addr : directNoReturnCalls)
+		{
+			noRets[i].arch = addr.arch->GetObject();
+			noRets[i].address = addr.address;
+			i++;
+		}
+
+		BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(m_context, noRets, directNoReturnCalls.size());
+		delete[] noRets;
+	}
+
+	if (m_haltedDisassemblyAddresses)
+	{
+		auto& haltedDisassemblyAddresses = *m_haltedDisassemblyAddresses;
+
+		BNArchitectureAndAddress* haltedAddresses = new BNArchitectureAndAddress[haltedDisassemblyAddresses.size()];
+
+		size_t i = 0;
+		for (auto& addr : haltedDisassemblyAddresses)
+		{
+			haltedAddresses[i].arch = addr.arch->GetObject();
+			haltedAddresses[i].address = addr.address;
+			i++;
+		}
+
+		BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(m_context, haltedAddresses, haltedDisassemblyAddresses.size());
+		delete[] haltedAddresses;
+	}
+
+	if (m_contextualReturns)
+	{
+		auto& contextualReturns = *m_contextualReturns;
+
+		bool dirty = contextualReturns.size() != m_context->contextualFunctionReturnCount;
+
+		if (!dirty)
+		{
+			size_t i = 0;
+			for (auto& pair : contextualReturns)
+			{
+				if (pair.first.arch->GetObject() != m_context->contextualFunctionReturnLocations[i].arch
+				 || pair.first.address != m_context->contextualFunctionReturnLocations[i].address
+				 || pair.second != m_context->contextualFunctionReturnValues[i])
+				{
+					dirty = true;
+					break;
+				}
+
+				i++;
+			}
+		}
+
+		if (dirty)
+		{
+			BNArchitectureAndAddress* returns = new BNArchitectureAndAddress[contextualReturns.size()];
+			bool* values = new bool[contextualReturns.size()];
+
+			size_t i = 0;
+			for (auto& pair : contextualReturns)
+			{
+				returns[i].arch = pair.first.arch->GetObject();
+				returns[i].address = pair.first.address;
+				values[i] = pair.second;
+				i++;
+			}
+
+			BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(m_context, returns, values, contextualReturns.size());
+
+			delete[] returns;
+			delete[] values;
+		}
+	}
+
+	BNAnalyzeBasicBlocksContextFinalize(m_context);
+}
+
 
 Architecture::Architecture(BNArchitecture* arch)
 {
@@ -332,8 +561,10 @@ void Architecture::AnalyzeBasicBlocksCallback(void *ctxt, BNFunction* function,
 	BNBasicBlockAnalysisContext* context)
 {
 	CallbackRef<Architecture> arch(ctxt);
+	BasicBlockAnalysisContext abbc(context);
 	Ref<Function> func(new Function(BNNewFunctionReference(function)));
-	arch->AnalyzeBasicBlocks(*func, *context);
+
+	arch->AnalyzeBasicBlocks(func, abbc);
 }
 
 
@@ -956,7 +1187,7 @@ bool Architecture::GetInstructionLowLevelIL(const uint8_t*, uint64_t, size_t&, L
 }
 
 
-void Architecture::AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context)
+void Architecture::AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context)
 {
 	DefaultAnalyzeBasicBlocks(function, context);
 }
@@ -1522,9 +1753,9 @@ bool CoreArchitecture::GetInstructionLowLevelIL(const uint8_t* data, uint64_t ad
 }
 
 
-void CoreArchitecture::AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context)
+void CoreArchitecture::AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context)
 {
-	BNArchitectureAnalyzeBasicBlocks(m_object, function.GetObject(), &context);
+	BNArchitectureAnalyzeBasicBlocks(m_object, function->GetObject(), context.m_context);
 }
 
 

binaryninjaapi.h

@@ -8073,7 +8073,52 @@ namespace BinaryNinja {
 	class RelocationHandler;
 
 	typedef size_t ExprId;
-	typedef BNBasicBlockAnalysisContext BasicBlockAnalysisContext;
+
+	class BasicBlockAnalysisContext
+	{
+	private:
+		// in
+		std::optional<std::map<ArchAndAddr, std::set<ArchAndAddr>>> m_indirectBranches;
+		std::optional<std::set<ArchAndAddr>> m_indirectNoReturnCalls;
+
+		// in/out
+		std::optional<std::map<ArchAndAddr, bool>> m_contextualReturns;
+
+		// out
+		std::optional<std::map<uint64_t, std::set<ArchAndAddr>>> m_directCodeReferences;
+		std::optional<std::set<ArchAndAddr>> m_directNoReturnCalls;
+		std::optional<std::set<ArchAndAddr>> m_haltedDisassemblyAddresses;
+
+	public:
+		BNBasicBlockAnalysisContext* m_context;
+
+		BasicBlockAnalysisContext(BNBasicBlockAnalysisContext* context);
+
+		BNFunctionAnalysisSkipOverride GetAnalysisSkipOverride() const { return m_context->analysisSkipOverride; }
+		bool GetTranslateTailCalls() const { return m_context->translateTailCalls; }
+		bool GetDisallowBranchToString() const { return m_context->disallowBranchToString; }
+		bool GetHaltOnInvalidInstructions() const { return m_context->haltOnInvalidInstructions; }
+		uint64_t GetMaxFunctionSize() const { return m_context->maxFunctionSize; }
+
+		bool GetMaxSizeReached() const { return m_context->maxSizeReached; }
+		void SetMaxSizeReached(bool reached) { m_context->maxSizeReached = reached; }
+
+		const std::map<ArchAndAddr, std::set<ArchAndAddr>> GetIndirectBranches();
+		const std::set<ArchAndAddr>& GetIndirectNoReturnCalls();
+
+		std::map<ArchAndAddr, bool>& GetContextualReturns();
+
+		std::map<uint64_t, std::set<ArchAndAddr>>& GetDirectCodeReferences();
+		std::set<ArchAndAddr>& GetDirectNoReturnCalls();
+		std::set<ArchAndAddr>& GetHaltedDisassemblyAddresses();
+
+		void AddTempOutgoingReference(Function* targetFunc);
+
+		Ref<BasicBlock> CreateBasicBlock(Architecture* arch, uint64_t start);
+		void AddFunctionBasicBlock(BasicBlock* block);
+
+		void Finalize();
+	};
 
 	/*! The Architecture class is the base class for all CPU architectures. This provides disassembly, assembly,
 	    patching, and IL translation lifting for a given architecture.
@@ -8182,7 +8227,7 @@ namespace BinaryNinja {
 			\param function Function to analyze
 			\param context Context for the analysis
 		*/
-		static void DefaultAnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context);
+		static void DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context);
 
 		/*! Get an Architecture by name
 
@@ -8287,7 +8332,7 @@ namespace BinaryNinja {
 			\param function Function to analyze
 			\param context Context for the analysis
 		*/
-		virtual void AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context);
+		virtual void AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context);
 
 		/*! Gets a register name from a register index.
 
@@ -8682,7 +8727,7 @@ namespace BinaryNinja {
 		    const uint8_t* data, uint64_t addr, size_t& len, std::vector<InstructionTextToken>& result) override;
 		virtual bool GetInstructionLowLevelIL(
 		    const uint8_t* data, uint64_t addr, size_t& len, LowLevelILFunction& il) override;
-		virtual void AnalyzeBasicBlocks(Function& function, BasicBlockAnalysisContext& context) override;
+		virtual void AnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context) override;
 		virtual std::string GetRegisterName(uint32_t reg) override;
 		virtual std::string GetFlagName(uint32_t flag) override;
 		virtual std::string GetFlagWriteTypeName(uint32_t flags) override;
@@ -11063,24 +11108,6 @@ namespace BinaryNinja {
 		*/
 		std::vector<Ref<BasicBlock>> GetBasicBlocks() const;
 
-		/*! Create a new basic block for this function
-
-			\param arch Architecture for the basic block
-			\param addr Address of the basic block
-			\return The new BasicBlock
-		*/
-		Ref<BasicBlock> CreateBasicBlock(Architecture* arch, uint64_t addr);
-
-		/*! Add a basic block to the function analysis basic block list
-
-			\param block The BasicBlock to add
-		*/
-		void AddBasicBlock(Ref<BasicBlock> block);
-
-		/*! Finalize basic block list for this function
-		*/
-		void FinalizeBasicBlocks();
-
 		/*! Get the basic block an address is located in
 
 			\param arch Architecture for the basic block
@@ -11411,16 +11438,7 @@ namespace BinaryNinja {
 		std::vector<IndirectBranchInfo> GetIndirectBranches();
 		std::vector<IndirectBranchInfo> GetIndirectBranchesAt(Architecture* arch, uint64_t addr);
 
-		void AddDirectCodeReference(const ArchAndAddr& source, uint64_t target);
-		void AddDirectNoReturnCall(const ArchAndAddr& location);
-		bool LocationHasNoReturnCalls(const ArchAndAddr& location) const;
 		Ref<Function> GetCalleeForAnalysis(Ref<Platform> platform, uint64_t addr, bool exact);
-		void AddTempOutgoingReference(Ref<Function> target);
-		bool HasTempOutgoingReference(Ref<Function> target) const;
-		void AddTempIncomingReference(Ref<Function> source);
-
-		bool GetContextualFunctionReturn(const ArchAndAddr& location, bool& value) const;
-		void SetContextualFunctionReturn(const ArchAndAddr& location, bool value);
 
 		std::vector<uint64_t> GetUnresolvedIndirectBranches();
 		bool HasUnresolvedIndirectBranches();