Implement EFI resolver as a module workflow

1 files changed, 274 insertions, 0 deletions

diff --git a/plugins/efi_resolver/src/DxeResolver.cpp b/plugins/efi_resolver/src/DxeResolver.cpp
new file mode 100644
index 00000000..f31c42ae
--- /dev/null
+++ b/plugins/efi_resolver/src/DxeResolver.cpp
@@ -0,0 +1,274 @@
+#include "DxeResolver.h"
+
+bool DxeResolver::resolveBootServices()
+{
+	m_task->SetProgressText("Resolving Boot Services...");
+	auto refs = m_view->GetCodeReferencesForType(QualifiedName("EFI_BOOT_SERVICES"));
+	// search reference of `EFI_BOOT_SERVICES` so that we can easily parse different services
+
+	for (auto& ref : refs)
+	{
+		if (m_task->IsCancelled())
+			return false;
+
+		auto func = ref.func;
+		auto mlil = func->GetMediumLevelIL();
+		if (!mlil)
+			continue;
+
+		auto mlilSsa = mlil->GetSSAForm();
+		size_t mlilIdx = mlil->GetInstructionStart(m_view->GetDefaultArchitecture(), ref.addr);
+		auto instr = mlilSsa->GetInstruction(mlil->GetSSAInstructionIndex(mlilIdx));
+
+		if (instr.operation == MLIL_CALL_SSA || instr.operation == MLIL_TAILCALL_SSA)
+		{
+			auto dest = instr.GetDestExpr();
+			if (dest.operation != MLIL_LOAD_STRUCT_SSA)
+				continue;
+			auto offset = dest.GetOffset();
+
+			if (offset == 0x18 + m_width * 16 || offset == 0x18 + m_width * 32)
+			{
+				// HandleProtocol, OpenProtocol
+				// Guid:1, Interface:2
+				resolveGuidInterface(ref.func, ref.addr, 1, 2);
+			}
+			else if (offset == 0x18 + m_width * 37)
+			{
+				// LocateProtocol
+				resolveGuidInterface(ref.func, ref.addr, 0, 2);
+			}
+		}
+	}
+	return true;
+}
+
+bool DxeResolver::resolveRuntimeServices()
+{
+	m_task->SetProgressText("Resolving Runtime Services...");
+	auto refs = m_view->GetCodeReferencesForType(QualifiedName("EFI_RUNTIME_SERVICES"));
+
+	for (auto& ref : refs)
+	{
+		if (m_task->IsCancelled())
+			return false;
+
+		auto func = ref.func;
+		auto mlil = func->GetMediumLevelIL();
+		if (!mlil)
+			continue;
+
+		auto mlilSsa = mlil->GetSSAForm();
+		size_t mlilIdx = mlil->GetInstructionStart(m_view->GetDefaultArchitecture(), ref.addr);
+		auto instr = mlilSsa->GetInstruction(mlil->GetSSAInstructionIndex(mlilIdx));
+
+		if (instr.operation == MLIL_CALL_SSA || instr.operation == MLIL_TAILCALL_SSA)
+		{
+			auto dest = instr.GetDestExpr();
+			if (dest.operation != MLIL_LOAD_STRUCT_SSA)
+				continue;
+			auto offset = dest.GetOffset();
+			if (offset == 0x18 + m_width * 6 || offset == 0x18 + m_width * 8)
+			{
+				// TODO implement this
+				// GetVariable and SetVariable
+			}
+		}
+	}
+	return true;
+}
+
+bool DxeResolver::resolveSmmTables(string serviceName, string tableName)
+{
+	m_task->SetProgressText("Defining MM tables...");
+	auto refs = m_view->GetCodeReferencesForType(QualifiedName(serviceName));
+	// both versions use the same type, so we only need to search for this one
+	for (auto& ref : refs)
+	{
+		if (m_task->IsCancelled())
+			return false;
+
+		auto func = ref.func;
+		auto mlil = func->GetMediumLevelIL();
+		if (!mlil)
+			continue;
+
+		auto mlilSsa = mlil->GetSSAForm();
+		size_t mlilIdx = mlil->GetInstructionStart(m_view->GetDefaultArchitecture(), ref.addr);
+		auto instr = mlilSsa->GetInstruction(mlil->GetSSAInstructionIndex(mlilIdx));
+
+		if (instr.operation != MLIL_CALL_SSA && instr.operation != MLIL_TAILCALL_SSA)
+			continue;
+
+		auto destExpr = instr.GetDestExpr();
+		if (destExpr.operation != MLIL_LOAD_STRUCT_SSA)
+			continue;
+
+		if (destExpr.GetOffset() != 8)
+			continue;
+
+		auto params = instr.GetParameterExprs();
+		if (params.size() < 2)
+			continue;
+
+		auto smstAddr = params[1];
+		if (smstAddr.operation != MLIL_CONST_PTR)
+			continue;
+
+		QualifiedNameAndType result;
+		string errors;
+		bool ok = m_view->ParseTypeString(tableName, result, errors);
+		if (!ok)
+			return false;
+		m_view->DefineDataVariable(smstAddr.GetValue().value, result.type);
+		m_view->DefineUserSymbol(new Symbol(DataSymbol, "gMmst", smstAddr.GetValue().value));
+		m_view->UpdateAnalysisAndWait();
+	}
+	return true;
+}
+
+bool DxeResolver::resolveSmmServices()
+{
+	m_task->SetProgressText("Resolving MM services...");
+	auto refs = m_view->GetCodeReferencesForType(QualifiedName("EFI_MM_SYSTEM_TABLE"));
+	auto refs_smm = m_view->GetCodeReferencesForType(QualifiedName("EFI_SMM_SYSTEM_TABLE2"));
+	// These tables have same type information, we can just iterate once
+	refs.insert(refs.end(), refs_smm.begin(), refs_smm.end());
+
+	for (auto& ref : refs)
+	{
+		if (m_task->IsCancelled())
+			return false;
+
+		auto func = ref.func;
+		auto mlil = func->GetMediumLevelIL();
+		if (!mlil)
+			continue;
+
+		auto mlilSsa = mlil->GetSSAForm();
+		size_t mlilIdx = mlil->GetInstructionStart(m_view->GetDefaultArchitecture(), ref.addr);
+		auto instr = mlilSsa->GetInstruction(mlil->GetSSAInstructionIndex(mlilIdx));
+
+		if (instr.operation == MLIL_CALL_SSA || instr.operation == MLIL_TAILCALL_SSA)
+		{
+			auto dest = instr.GetDestExpr();
+			if (dest.operation != MLIL_LOAD_STRUCT_SSA)
+				continue;
+			auto offset = dest.GetOffset();
+
+			if (offset == 0x18 + m_width * 0x14)
+			{
+				// SmmHandleProtocol
+				resolveGuidInterface(ref.func, ref.addr, 1, 2);
+			}
+			else if (offset == 0x18 + m_width * 0x17)
+			{
+				// SmmLocateProtocol
+				resolveGuidInterface(ref.func, ref.addr, 0, 2);
+			}
+		}
+	}
+	return true;
+}
+
+bool DxeResolver::resolveSmiHandlers()
+{
+	m_task->SetProgressText("Resolving SMI Handlers...");
+	auto refs = m_view->GetCodeReferencesForType(QualifiedName("EFI_MM_SW_REGISTER"));
+	auto refs_smm_sw = m_view->GetCodeReferencesForType(QualifiedName("EFI_SMM_SW_REGISTER2"));
+	auto refs_mm_sx = m_view->GetCodeReferencesForType(QualifiedName("EFI_MM_SX_REGISTER"));
+	auto refs_smm_sx = m_view->GetCodeReferencesForType(QualifiedName("EFI_SMM_SX_REGISTER2"));
+	// Define them together
+
+	refs.insert(refs.end(), refs_smm_sw.begin(), refs_smm_sw.end());
+	refs.insert(refs.end(), refs_smm_sx.begin(), refs_smm_sw.end());
+	refs.insert(refs.end(), refs_mm_sx.begin(), refs_mm_sx.end());
+
+	for (auto& ref : refs)
+	{
+		if (m_task->IsCancelled())
+			return false;
+
+		auto func = ref.func;
+		auto mlil = func->GetMediumLevelIL();
+		if (!mlil)
+			continue;
+
+		auto mlilSsa = mlil->GetSSAForm();
+		size_t mlilIdx = mlil->GetInstructionStart(m_view->GetDefaultArchitecture(), ref.addr);
+		auto instr = mlilSsa->GetInstruction(mlil->GetSSAInstructionIndex(mlilIdx));
+
+		if (instr.operation == MLIL_CALL_SSA || instr.operation == MLIL_TAILCALL_SSA)
+		{
+			auto dest = instr.GetDestExpr();
+			if (dest.operation != MLIL_LOAD_STRUCT_SSA)
+				continue;
+
+			auto offset = dest.GetOffset();
+			if (offset == 0)
+			{
+				auto parameters = instr.GetParameterExprs();
+				if (parameters.size() < 4)
+					continue;
+
+				// TODO we should be able to parse registerContext, but it's normally an aliased variable
+				//    and we have some issues relate to that
+				auto dispatchFunction = parameters[1];
+				if (dispatchFunction.operation != MLIL_CONST_PTR)
+					continue;
+				auto funcAddr = static_cast<uint64_t>(dispatchFunction.GetConstant());
+				auto targetFunc = m_view->GetAnalysisFunction(m_view->GetDefaultPlatform(), funcAddr);
+				auto funcType = targetFunc->GetType();
+				std::ostringstream ss;
+				ss << "SmiHandler_" << std::hex << funcAddr;
+				string funcName = ss.str();
+
+				// typedef enum
+				string handleTypeStr =
+					"EFI_STATUS SmiHandler(EFI_HANDLE DispatchHandle, VOID* Context, VOID* CommBuffer, UINTN* "
+					"CommBufferSize);";
+				QualifiedNameAndType result;
+				string errors;
+				bool ok = m_view->ParseTypeString(handleTypeStr, result, errors);
+				if (!ok)
+					return false;
+				targetFunc->SetUserType(result.type);
+				m_view->DefineUserSymbol(new Symbol(FunctionSymbol, funcName, funcAddr));
+				m_view->UpdateAnalysisAndWait();
+
+				// After setting the type, we want to propagate the parameters' type
+				TypePropagation propagator(m_view);
+				propagator.propagateFuncParamTypes(targetFunc);
+			}
+		}
+	}
+	return true;
+}
+
+bool DxeResolver::resolveDxe()
+{
+	if (!resolveBootServices())
+		return false;
+	if (!resolveRuntimeServices())
+		return false;
+	return true;
+}
+
+bool DxeResolver::resolveSmm()
+{
+	if (!resolveSmmTables("EFI_SMM_GET_SMST_LOCATION2", "EFI_SMM_SYSTEM_TABLE2*"))
+		return false;
+	if (!resolveSmmTables("EFI_MM_GET_MMST_LOCATION", "EFI_MM_SYSTEM_TABLE*"))
+		return false;
+	if (!resolveSmmServices())
+		return false;
+	if (!resolveSmiHandlers())
+		return false;
+	return true;
+}
+
+DxeResolver::DxeResolver(Ref<BinaryView> view, Ref<BackgroundTask> task) : Resolver(view, task)
+{
+	initProtocolMapping();
+	setModuleEntry(DXE);
+}