1 files changed, 51 insertions, 31 deletions

diff --git a/view/sharedcache/workflow/SharedCacheWorkflow.cpp b/view/sharedcache/workflow/SharedCacheWorkflow.cpp
index 6aab9e50..5980d141 100644
--- a/view/sharedcache/workflow/SharedCacheWorkflow.cpp
+++ b/view/sharedcache/workflow/SharedCacheWorkflow.cpp
@@ -159,13 +159,13 @@ void AnalyzeStubFunction(Ref<Function> func, Ref<MediumLevelILFunction> mlil, Sh
 				if (defInstr.operation != MLIL_SET_VAR_SSA)
 					return;
 				expr = defInstr.GetSourceExpr<MLIL_SET_VAR_SSA>();
-				// Fallthrough to load ptr.
+				// Fallthrough to MLIL_LOAD_SSA.
 			}
 		case MLIL_LOAD_SSA:
 			expr = expr.GetSourceExpr<MLIL_LOAD_SSA>();
 			if (expr.operation != MLIL_CONST_PTR)
 				return;
-			// Fallthrough to const ptr.
+			// Fallthrough to MLIL_CONST_PTR.
 		case MLIL_CONST_PTR:
 			{
 				// First load the stub island, if we _do_ load the stub island stop and reanalyze for constant propagation.'
@@ -183,7 +183,7 @@ void AnalyzeStubFunction(Ref<Function> func, Ref<MediumLevelILFunction> mlil, Sh
 		}
 	};
 
-	auto processTailcallExpr = [&](MediumLevelILInstruction expr) {
+	auto processTailcallExpr = [&](const MediumLevelILInstruction& expr) {
 		switch (expr.operation)
 		{
 		case MLIL_CONST_PTR:
@@ -222,7 +222,7 @@ void AnalyzeStandardFunction(Ref<Function> func, Ref<MediumLevelILFunction> mlil
 	auto view = func->GetView();
 	auto identifyUnmappedSymbol = [&](uint64_t symbolAddr) {
 		// Skip if already loaded.
-		if (view->IsValidOffset(symbolAddr))
+		if (view->IsValidOffset(symbolAddr) || view->GetSymbolByAddress(symbolAddr))
 			return false;
 		const auto symbol = controller.GetSymbolAt(symbolAddr);
 		if (!symbol.has_value())
@@ -235,52 +235,72 @@ void AnalyzeStandardFunction(Ref<Function> func, Ref<MediumLevelILFunction> mlil
 		// Skip if already loaded.
 		if (view->IsValidOffset(regionAddr))
 			return false;
-		const auto region = controller.GetRegionContaining(regionAddr);
+		auto region = controller.GetRegionContaining(regionAddr);
 		if (!region.has_value())
 			return false;
-		// Only interested in stub islands which would contain the pointers to other image functions.
-		if (region->type != SharedCacheRegionTypeStubIsland)
+		// Only interested in non image regions, we DON'T want to implicitly load image regions (with functions presumably).
+		if (region->type == SharedCacheRegionTypeImage)
 			return false;
-
+		// Adjust the new region semantics to read only, this helps analysis pickup constant loads in our stub functions.
+		// NOTE: We do NOT do this for stub island as that contains CODE!
+		if (region->type != SharedCacheRegionTypeStubIsland)
+			region->flags = static_cast<BNSegmentFlag>(SegmentReadable | SegmentContainsData);
 		return controller.ApplyRegion(*view, *region);
 	};
 
-	auto processJumpExpr = [&](const MediumLevelILInstruction& expr) {
+	// Use this to filter out constants that are impossible to be backed by a region.
+	auto baseAddress = view->GetStart();
+
+	// Promotes a constant to a constant pointer, if it's backed by a shared cache region.
+	// Constants won't be eligible for symbol rendering so to get it to render we must rewrite the IL to make
+	// it a constant pointer.
+	auto promoteUnmappedPointerExpr = [&](MediumLevelILInstruction constExpr) {
+		const auto unmappedAddr = constExpr.GetConstant<MLIL_CONST>();
+		if (unmappedAddr < baseAddress || view->IsValidOffset(unmappedAddr))
+			return false;
+		if (!controller.GetRegionContaining(unmappedAddr))
+			return false;
+		// Replace the constant expr with a constant pointer expr
+		const auto constSrcLoc = ILSourceLocation(constExpr.address, constExpr.sourceOperand);
+		const auto constPtrExpr = mlil->ConstPointer(constExpr.size, unmappedAddr, constSrcLoc);
+		// Force the expr type to a pointer as well
+		const auto ptrType = Type::PointerType(func->GetArchitecture(), Type::VoidType()->WithConfidence(0))->WithConfidence(BN_HEURISTIC_CONFIDENCE);
+		mlil->SetExprType(constPtrExpr, ptrType);
+		constExpr.Replace(constPtrExpr);
+		return true;
+	};
+
+	auto processUnmappedExpr = [&](const MediumLevelILInstruction& expr) {
 		switch (expr.operation)
 		{
 		case MLIL_CONST_PTR:
 			loadStubIslandRegion(expr.GetConstant<MLIL_CONST_PTR>());
 			identifyUnmappedSymbol(expr.GetConstant<MLIL_CONST_PTR>());
 			break;
+		case MLIL_CONST:
+			// Typically a direct expression load/store will have the constant be promoted to a constant pointer
+			// however if an expression is only used as an address of and data flow does not fold it in, we won't get a
+			// constant pointer promotion. This case handles that limitation, by checking to see if a constant falls in
+			// an unmapped region.
+			if (promoteUnmappedPointerExpr(expr))
+			{
+				// We have promoted a pointer, lets now try and identify its symbol.
+				// We don't want to rely on this activity being re-invoked and going down the `MLIL_CONST_PTR` path
+				// so instead we just call the `identifyUnmappedSymbol` after promotion, to ensure that the symbols are identified.
+				identifyUnmappedSymbol(expr.GetConstant());
+			}
+			break;
 		default:
 			break;
 		}
+		return true;
 	};
 
-	// Load all unmapped STUB regions / images that are called in this function.
+	// 1. Load all unmapped STUB regions / images that are called in this function.
+	// 2. Identify loads & stores to unmapped regions and add their respective symbol.
 	for (const auto& block : mlil->GetBasicBlocks())
-	{
 		for (size_t i = block->GetStart(), end = block->GetEnd(); i < end; ++i)
-		{
-			auto instr = mlil->GetInstruction(i);
-			switch (instr.operation)
-			{
-				case MLIL_CALL_SSA:
-					processJumpExpr(instr.GetDestExpr<MLIL_CALL_SSA>());
-					break;
-				case MLIL_JUMP:
-					processJumpExpr(instr.GetDestExpr<MLIL_JUMP>());
-					break;
-				default:
-					break;
-			}
-			// TODO: Check all instructions for accesses to select region types (stub etc...)
-			// TODO: ^ we actually dont really need to do this, the other type of access cont..
-			// TODO: the other two types of accesses (load & save) we dont want to load their regions, just
-			// TODO: their symbol information if available.
-			// TODO: See:
-		}
-	}
+			mlil->GetInstruction(i).VisitExprs(processUnmappedExpr);
 }
 
 void AnalyzeFunction(Ref<AnalysisContext> ctx)