Add UI action and API to control early returns in HLIL

Author: D0ntPanic

binaryninjaapi.h

@@ -11433,6 +11433,9 @@ namespace BinaryNinja {
 		bool IsConditionInverted(uint64_t addr);
 		void SetConditionInverted(uint64_t addr, bool invert);
 
+		BNEarlyReturn GetEarlyReturn(uint64_t addr);
+		void SetEarlyReturn(uint64_t addr, BNEarlyReturn mode);
+
 		std::map<Variable, std::set<Variable>> GetMergedVariables();
 		void MergeVariables(const Variable& target, const std::set<Variable>& sources);
 		void UnmergeVariables(const Variable& target, const std::set<Variable>& sources);

binaryninjacore.h

@@ -37,7 +37,7 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 109
+#define BN_CURRENT_CORE_ABI_VERSION 110
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
@@ -1027,6 +1027,9 @@ extern "C"
 
 		// HLIL condition can be displayed as the inverse
 		HLILInvertableCondition = 0x200,
+
+		// HLIL condition can be rewritten as an early return
+		HLILEarlyReturnPossible = 0x400,
 	} BNILInstructionAttribute;
 
 	typedef enum BNIntrinsicClass
@@ -3241,6 +3244,15 @@ extern "C"
 		AllowExprFolding
 	} BNExprFolding;
 
+	typedef enum BNEarlyReturn
+	{
+		DefaultEarlyReturn,
+		PreventEarlyReturn,
+		SmallestSideEarlyReturn,
+		TrueSideEarlyReturn,
+		FalseSideEarlyReturn
+	} BNEarlyReturn;
+
 	typedef struct BNDebugFunctionInfo
 	{
 		char* shortName;
@@ -4998,6 +5010,8 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetExprFolding(BNFunction* func, uint64_t addr, BNExprFolding mode);
 	BINARYNINJACOREAPI bool BNIsConditionInverted(BNFunction* func, uint64_t addr);
 	BINARYNINJACOREAPI void BNSetConditionInverted(BNFunction* func, uint64_t addr, bool invert);
+	BINARYNINJACOREAPI BNEarlyReturn BNGetEarlyReturn(BNFunction* func, uint64_t addr);
+	BINARYNINJACOREAPI void BNSetEarlyReturn(BNFunction* func, uint64_t addr, BNEarlyReturn mode);
 	BINARYNINJACOREAPI BNMergedVariable* BNGetMergedVariables(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI void BNFreeMergedVariableList(BNMergedVariable* vars, size_t count);
 	BINARYNINJACOREAPI void BNMergeVariables(BNFunction* func, const BNVariable* target, const BNVariable* sources,

function.cpp

@@ -2837,6 +2837,18 @@ void Function::SetConditionInverted(uint64_t addr, bool invert)
 }
 
 
+BNEarlyReturn Function::GetEarlyReturn(uint64_t addr)
+{
+	return BNGetEarlyReturn(m_object, addr);
+}
+
+
+void Function::SetEarlyReturn(uint64_t addr, BNEarlyReturn mode)
+{
+	BNSetEarlyReturn(m_object, addr, mode);
+}
+
+
 std::map<Variable, std::set<Variable>> Function::GetMergedVariables()
 {
 	size_t count;

python/function.py

@@ -29,7 +29,7 @@
 from .enums import (
 	AnalysisSkipReason, FunctionGraphType, SymbolType, InstructionTextTokenType, HighlightStandardColor,
 	HighlightColorStyle, DisassemblyOption, IntegerDisplayType, FunctionAnalysisSkipOverride, FunctionUpdateType,
-	BuiltinType, ExprFolding
+	BuiltinType, ExprFolding, EarlyReturn
 )
 
 from . import associateddatastore  # Required in the main scope due to being an argument for _FunctionAssociatedDataStore
@@ -3476,6 +3476,16 @@ def set_condition_inverted(self, addr: Union[int, highlevelil.HighLevelILInstruc
 			addr = addr.address
 		core.BNSetConditionInverted(self.handle, addr, invert)
 
+	def get_early_return(self, addr: Union[int, highlevelil.HighLevelILInstruction]) -> EarlyReturn:
+		if isinstance(addr, highlevelil.HighLevelILInstruction):
+			addr = addr.address
+		return EarlyReturn(core.BNGetEarlyReturn(self.handle, addr))
+
+	def set_early_return(self, addr: Union[int, highlevelil.HighLevelILInstruction], value: EarlyReturn):
+		if isinstance(addr, highlevelil.HighLevelILInstruction):
+			addr = addr.address
+		core.BNSetEarlyReturn(self.handle, addr, value)
+
 
 class AdvancedFunctionAnalysisDataRequestor:
 	def __init__(self, func: Optional['Function'] = None):

ui/commands.h

@@ -98,6 +98,7 @@ TypeRef GetFunctionType(BinaryViewRef data, TypeRef type);
 std::optional<uint64_t> getFoldableExprAddress(
 	BinaryNinja::HighLevelILFunction* hlil, const HighlightTokenState& highlight);
 std::optional<uint64_t> getInvertableConditionAddress(BinaryNinja::HighLevelILFunction* hlil, size_t instrIndex);
+std::optional<uint64_t> getEarlyReturnAddress(BinaryNinja::HighLevelILFunction* hlil, size_t instrIndex);
 
 /*!
     @}

ui/flowgraphwidget.h

@@ -170,6 +170,8 @@ class BINARYNINJAUIAPI FlowGraphWidget :
 	BNExprFolding getCurrentExprFolding();
 	std::optional<uint64_t> getCurrentInvertableConditionAddress();
 	bool getCurrentConditionInverted();
+	std::optional<uint64_t> getCurrentEarlyReturnAddress();
+	BNEarlyReturn getCurrentEarlyReturn();
 	std::optional<std::pair<BinaryNinja::Variable, BinaryNinja::Variable>> getMergeVariablesAtCurrentLocation();
 
   protected:
@@ -401,6 +403,7 @@ class BINARYNINJAUIAPI FlowGraphWidget :
 	void setCurrentVariableDeadStoreElimination(BNDeadStoreElimination elimination);
 	void setCurrentExprFolding(BNExprFolding folding);
 	void toggleConditionInverted();
+	void setCurrentEarlyReturn(BNEarlyReturn earlyReturn);
 	void splitToNewTabAndNavigateFromCursorPosition();
 	void splitToNewWindowAndNavigateFromCursorPosition();
 	void splitToNewPaneAndNavigateFromCursorPosition();

ui/linearview.h

@@ -320,6 +320,8 @@ class BINARYNINJAUIAPI LinearView : public QAbstractScrollArea, public View, pub
 	BNExprFolding getCurrentExprFolding();
 	std::optional<uint64_t> getCurrentInvertableConditionAddress();
 	bool getCurrentConditionInverted();
+	std::optional<uint64_t> getCurrentEarlyReturnAddress();
+	BNEarlyReturn getCurrentEarlyReturn();
 
 	void setDataButtonVisible(bool visible);
 	std::optional<std::pair<BinaryNinja::Variable, BinaryNinja::Variable>> getMergeVariablesAtCurrentLocation();
@@ -412,6 +414,7 @@ private Q_SLOTS:
 	void setCurrentVariableDeadStoreElimination(BNDeadStoreElimination elimination);
 	void setCurrentExprFolding(BNExprFolding folding);
 	void toggleConditionInverted();
+	void setCurrentEarlyReturn(BNEarlyReturn earlyReturn);
 
 Q_SIGNALS:
 	void notifyResizeEvent(int width, int height);