From 84124764143d4f35961e2c61494d97640ac7f99f Mon Sep 17 00:00:00 2001
From: Brian Potchik <brian@vector35.com>
Date: Wed, 25 Aug 2021 22:59:50 -0400
Subject: Update workflows c++ examples.

---
 examples/workflows/inliner/inliner.cpp       | 23 ++++++++++++++-----
 examples/workflows/objectivec/objectivec.cpp | 33 +++++++++++++++++-----------
 2 files changed, 37 insertions(+), 19 deletions(-)

(limited to 'examples/workflows')

diff --git a/examples/workflows/inliner/inliner.cpp b/examples/workflows/inliner/inliner.cpp
index 01eacbd2..4ce05a31 100644
--- a/examples/workflows/inliner/inliner.cpp
+++ b/examples/workflows/inliner/inliner.cpp
@@ -4,6 +4,7 @@
 #include <cinttypes>
 #include <cstdint>
 #include <cstdio>
+#include <mutex>
 #include <string>
 #include <tuple>
 #include <unordered_map>
@@ -25,16 +26,26 @@ extern "C"
 {
 	BN_DECLARE_CORE_ABI_VERSION
 
-	unordered_map<uint64_t, set<uint64_t>> g_callSiteInlines;
+	// TODO: Replace with analysis cache opaque datastore.
+	std::mutex g_mutex;
+	unordered_map<void*, unordered_map<uint64_t, set<uint64_t>>> g_callSiteInlines;
 
 	void FunctionInliner(Ref<AnalysisContext> analysisContext)
 	{
+		std::unique_lock<std::mutex> lock(g_mutex);
+
 		Ref<Function> function = analysisContext->GetFunction();
 		Ref<BinaryView> data = function->GetView();
+		auto gItr = g_callSiteInlines.find(data->GetObject());
+		if (gItr == g_callSiteInlines.end())
+			return;
+
+		auto itr = gItr->second.find(function->GetStart());
+		if (itr == gItr->second.end())
+			return;
 
-		set<uint64_t> callSiteInlines;
-		if (const auto& itr = g_callSiteInlines.find(function->GetStart()); itr != g_callSiteInlines.end())
-			callSiteInlines = itr->second;
+		auto& callSiteInlines = itr->second;
+		lock.unlock();
 
 		bool updated = false;
 		uint8_t opcode[BN_MAX_INSTRUCTION_LENGTH];
@@ -95,7 +106,6 @@ extern "C"
 								{
 									LowLevelILLabel label;
 									label.operand = instrIndex + 1;
-									llilFunc->AddInstruction(llilFunc->Pop(instr.size));
 									llilFunc->AddInstruction(llilFunc->Goto(label));
 								}
 								else
@@ -140,7 +150,8 @@ extern "C"
 			[](BinaryView* view, Function* func) {
 				// TODO func->Inform("inlinedCallSites")
 				// TODO resolve multiple embedded inlines
-				g_callSiteInlines[func->GetStart()].insert(view->GetCurrentOffset());
+				std::lock_guard<std::mutex> lock(g_mutex);
+				g_callSiteInlines[view->GetObject()][func->GetStart()].insert(view->GetCurrentOffset());
 				func->Reanalyze();
 			}, inlinerIsValid);
 
diff --git a/examples/workflows/objectivec/objectivec.cpp b/examples/workflows/objectivec/objectivec.cpp
index e967c5e3..f6cc4e47 100644
--- a/examples/workflows/objectivec/objectivec.cpp
+++ b/examples/workflows/objectivec/objectivec.cpp
@@ -25,6 +25,10 @@ extern "C"
 {
 	BN_DECLARE_CORE_ABI_VERSION
 
+	// TODO: Replace with analysis cache opaque datastore.
+	std::mutex g_mutex;
+	unordered_map<void*, unordered_map<uint64_t, tuple<string, uint64_t>>> g_classData;
+
 	// TODO
 	// * __objc_const missing xrefs to implementations
 	void ObjectiveCAnalysis(Ref<AnalysisContext> analysisContext)
@@ -32,9 +36,10 @@ extern "C"
 		Ref<Function> function = analysisContext->GetFunction();
 		Ref<BinaryView> data = function->GetView();
 
-		static unordered_map<uint64_t, tuple<string, uint64_t>> classData;
-		static bool sInit = false;
-		if (!sInit) // TODO move to module-based workflow activity with run-once semantics
+		// TODO 1) Move this to run-once workflow activity 2) stash data in analysis cache opaque datastore.
+		std::unique_lock<std::mutex> lock(g_mutex);
+		auto gItr = g_classData.find(data->GetObject());
+		if (gItr == g_classData.end())
 		{
 			auto constSection = data->GetSectionByName("__objc_const");
 			if (!constSection)
@@ -64,25 +69,27 @@ extern "C"
 				uint64_t impPtr = reader.Read64();
 				//string methodName = reader.ReadCString(selector);
 				string typeEncoding = "";//reader.ReadCString(typePtr);
-				classData.insert_or_assign(selector, std::forward_as_tuple(typeEncoding, impPtr));
+				g_classData[data->GetObject()].insert_or_assign(selector, std::forward_as_tuple(typeEncoding, impPtr));
 			}
 
-			reader.Seek(namePtr);
-			string className = reader.ReadCString();
-
-			auto cfStringSection = data->GetSectionByName("__cfstring");
-			if (!cfStringSection)
-				return;
+			// TODO cfstring
+			// reader.Seek(namePtr);
+			// string className = reader.ReadCString();
 
-			sInit = true;
+			// auto cfStringSection = data->GetSectionByName("__cfstring");
+			// if (!cfStringSection)
+			// 	return;
 		}
 
+		auto& classData = gItr->second;
+		lock.unlock();
+
 		bool updated = false;
 		uint8_t opcode[BN_MAX_INSTRUCTION_LENGTH];
 		InstructionInfo iInfo;
 
-		// // if (m_owner->IsAborted())
-		// // 	return;
+		// if (m_owner->IsAborted())
+		// 	return;
 
 		// TODO fix this....
 		auto sym = data->GetSymbolByRawName("_objc_msgSend");
-- 
cgit v1.3.1