Updating APIs to deal with stack adjustment

7 files changed, 84 insertions, 8 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 5fb95021..1ac60625 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1872,6 +1872,7 @@ namespace BinaryNinja
 		void SetConst(const Confidence<bool>& cnst);
 		void SetVolatile(const Confidence<bool>& vltl);
 		void SetTypeName(const QualifiedName& name);
+		Confidence<size_t> GetStackAdjustment() const;
 
 		uint64_t GetElementCount() const;
 		uint64_t GetOffset() const;
@@ -1911,7 +1912,8 @@ namespace BinaryNinja
 		static Ref<Type> ArrayType(const Confidence<Ref<Type>>& type, uint64_t elem);
 		static Ref<Type> FunctionType(const Confidence<Ref<Type>>& returnValue,
 			const Confidence<Ref<CallingConvention>>& callingConvention,
-			const std::vector<FunctionParameter>& params, const Confidence<bool>& varArg = Confidence<bool>(false, 0));
+			const std::vector<FunctionParameter>& params, const Confidence<bool>& varArg = Confidence<bool>(false, 0),
+			const Confidence<size_t>& stackAdjust = Confidence<size_t>(0, 0));
 
  		static std::string GenerateAutoTypeId(const std::string& source, const QualifiedName& name);
 		static std::string GenerateAutoDemangledTypeId(const QualifiedName& name);
@@ -2517,6 +2519,7 @@ namespace BinaryNinja
 		ExprId JumpTo(ExprId dest, const std::vector<BNLowLevelILLabel*>& targets,
 			const ILSourceLocation& loc = ILSourceLocation());
 		ExprId Call(ExprId dest, const ILSourceLocation& loc = ILSourceLocation());
+		ExprId CallStackAdjust(ExprId dest, size_t adjust, const ILSourceLocation& loc = ILSourceLocation());
 		ExprId CallSSA(const std::vector<SSARegister>& output, ExprId dest, const std::vector<SSARegister>& params,
 			const SSARegister& stack, size_t newMemoryVer, size_t prevMemoryVer,
 			const ILSourceLocation& loc = ILSourceLocation());
diff --git a/binaryninjacore.h b/binaryninjacore.h
index f20cc1f2..1bb4c726 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -323,6 +323,7 @@ extern "C"
 		LLIL_JUMP,
 		LLIL_JUMP_TO,
 		LLIL_CALL,
+		LLIL_CALL_STACK_ADJUST,
 		LLIL_RET,
 		LLIL_NORET,
 		LLIL_IF,
@@ -2655,7 +2656,7 @@ extern "C"
 	BINARYNINJACOREAPI BNType* BNCreateArrayType(BNTypeWithConfidence* type, uint64_t elem);
 	BINARYNINJACOREAPI BNType* BNCreateFunctionType(BNTypeWithConfidence* returnValue,
 		BNCallingConventionWithConfidence* callingConvention, BNFunctionParameter* params,
-		size_t paramCount, BNBoolWithConfidence* varArg);
+		size_t paramCount, BNBoolWithConfidence* varArg, BNSizeWithConfidence* stackAdjust);
 	BINARYNINJACOREAPI BNType* BNNewTypeReference(BNType* type);
 	BINARYNINJACOREAPI BNType* BNDuplicateType(BNType* type);
 	BINARYNINJACOREAPI char* BNGetTypeAndName(BNType* type, BNQualifiedName* name);
@@ -2688,6 +2689,7 @@ extern "C"
 	BINARYNINJACOREAPI void BNTypeSetMemberAccess(BNType* type, BNMemberAccessWithConfidence* access);
 	BINARYNINJACOREAPI void BNTypeSetConst(BNType* type, BNBoolWithConfidence* cnst);
 	BINARYNINJACOREAPI void BNTypeSetVolatile(BNType* type, BNBoolWithConfidence* vltl);
+	BINARYNINJACOREAPI BNSizeWithConfidence BNGetTypeStackAdjustment(BNType* type);
 
 	BINARYNINJACOREAPI char* BNGetTypeString(BNType* type, BNPlatform* platform);
 	BINARYNINJACOREAPI char* BNGetTypeStringBeforeName(BNType* type, BNPlatform* platform);
diff --git a/lowlevelilinstruction.cpp b/lowlevelilinstruction.cpp
index 19bfd983..d85e4f17 100644
--- a/lowlevelilinstruction.cpp
+++ b/lowlevelilinstruction.cpp
@@ -60,6 +60,7 @@ unordered_map<LowLevelILOperandUsage, LowLevelILOperandType>
 		{LowSSARegisterLowLevelOperandUsage, SSARegisterLowLevelOperand},
 		{ConstantLowLevelOperandUsage, IntegerLowLevelOperand},
 		{VectorLowLevelOperandUsage, IntegerLowLevelOperand},
+		{StackAdjustmentLowLevelOperandUsage, IntegerLowLevelOperand},
 		{TargetLowLevelOperandUsage, IndexLowLevelOperand},
 		{TrueTargetLowLevelOperandUsage, IndexLowLevelOperand},
 		{FalseTargetLowLevelOperandUsage, IndexLowLevelOperand},
@@ -111,6 +112,7 @@ unordered_map<BNLowLevelILOperation, vector<LowLevelILOperandUsage>>
 		{LLIL_JUMP, {DestExprLowLevelOperandUsage}},
 		{LLIL_JUMP_TO, {DestExprLowLevelOperandUsage, TargetListLowLevelOperandUsage}},
 		{LLIL_CALL, {DestExprLowLevelOperandUsage}},
+		{LLIL_CALL_STACK_ADJUST, {DestExprLowLevelOperandUsage, StackAdjustmentLowLevelOperandUsage}},
 		{LLIL_RET, {DestExprLowLevelOperandUsage}},
 		{LLIL_IF, {ConditionExprLowLevelOperandUsage, TrueTargetLowLevelOperandUsage,
 			FalseTargetLowLevelOperandUsage}},
@@ -1020,7 +1022,7 @@ LowLevelILInstruction LowLevelILInstructionBase::GetSSAForm() const
 		return *this;
 	size_t expr = GetSSAExprIndex();
 	size_t instr = GetSSAInstructionIndex();
-	return LowLevelILInstruction(ssa, ssa->GetRawExpr(GetSSAExprIndex()), expr, instr);
+	return LowLevelILInstruction(ssa, ssa->GetRawExpr(expr), expr, instr);
 }
 
 
@@ -1031,7 +1033,7 @@ LowLevelILInstruction LowLevelILInstructionBase::GetNonSSAForm() const
 		return *this;
 	size_t expr = GetNonSSAExprIndex();
 	size_t instr = GetNonSSAInstructionIndex();
-	return LowLevelILInstruction(nonSsa, nonSsa->GetRawExpr(GetSSAExprIndex()), expr, instr);
+	return LowLevelILInstruction(nonSsa, nonSsa->GetRawExpr(expr), expr, instr);
 }
 
 
@@ -1160,6 +1162,9 @@ void LowLevelILInstruction::VisitExprs(const std::function<bool(const LowLevelIL
 	case LLIL_CALL:
 		GetDestExpr<LLIL_CALL>().VisitExprs(func);
 		break;
+	case LLIL_CALL_STACK_ADJUST:
+		GetDestExpr<LLIL_CALL_STACK_ADJUST>().VisitExprs(func);
+		break;
 	case LLIL_CALL_SSA:
 		GetDestExpr<LLIL_CALL_SSA>().VisitExprs(func);
 		break;
@@ -1301,6 +1306,9 @@ ExprId LowLevelILInstruction::CopyTo(LowLevelILFunction* dest,
 		return dest->Jump(subExprHandler(GetDestExpr<LLIL_JUMP>()), *this);
 	case LLIL_CALL:
 		return dest->Call(subExprHandler(GetDestExpr<LLIL_CALL>()), *this);
+	case LLIL_CALL_STACK_ADJUST:
+		return dest->CallStackAdjust(subExprHandler(GetDestExpr<LLIL_CALL_STACK_ADJUST>()),
+			GetStackAdjustment<LLIL_CALL_STACK_ADJUST>(), *this);
 	case LLIL_RET:
 		return dest->Return(subExprHandler(GetDestExpr<LLIL_RET>()), *this);
 	case LLIL_JUMP_TO:
@@ -1647,6 +1655,15 @@ int64_t LowLevelILInstruction::GetVector() const
 }
 
 
+size_t LowLevelILInstruction::GetStackAdjustment() const
+{
+	size_t operandIndex;
+	if (GetOperandIndexForUsage(StackAdjustmentLowLevelOperandUsage, operandIndex))
+		return (size_t)GetRawOperandAsInteger(operandIndex);
+	throw LowLevelILInstructionAccessException();
+}
+
+
 size_t LowLevelILInstruction::GetTarget() const
 {
 	size_t operandIndex;
@@ -2133,6 +2150,12 @@ ExprId LowLevelILFunction::Call(ExprId dest, const ILSourceLocation& loc)
 }
 
 
+ExprId LowLevelILFunction::CallStackAdjust(ExprId dest, size_t adjust, const ILSourceLocation& loc)
+{
+	return AddExprWithLocation(LLIL_CALL_STACK_ADJUST, loc, 0, 0, dest, adjust);
+}
+
+
 ExprId LowLevelILFunction::CallSSA(const vector<SSARegister>& output, ExprId dest, const vector<SSARegister>& params,
 	const SSARegister& stack, size_t newMemoryVer, size_t prevMemoryVer, const ILSourceLocation& loc)
 {
diff --git a/lowlevelilinstruction.h b/lowlevelilinstruction.h
index 03688700..b2849d44 100644
--- a/lowlevelilinstruction.h
+++ b/lowlevelilinstruction.h
@@ -127,6 +127,7 @@ namespace BinaryNinja
 		LowSSARegisterLowLevelOperandUsage,
 		ConstantLowLevelOperandUsage,
 		VectorLowLevelOperandUsage,
+		StackAdjustmentLowLevelOperandUsage,
 		TargetLowLevelOperandUsage,
 		TrueTargetLowLevelOperandUsage,
 		FalseTargetLowLevelOperandUsage,
@@ -499,6 +500,7 @@ namespace BinaryNinja
 		template <BNLowLevelILOperation N> SSARegister GetLowSSARegister() const { return As<N>().GetLowSSARegister(); }
 		template <BNLowLevelILOperation N> int64_t GetConstant() const { return As<N>().GetConstant(); }
 		template <BNLowLevelILOperation N> int64_t GetVector() const { return As<N>().GetVector(); }
+		template <BNLowLevelILOperation N> size_t GetStackAdjustment() const { return As<N>().GetStackAdjustment(); }
 		template <BNLowLevelILOperation N> size_t GetTarget() const { return As<N>().GetTarget(); }
 		template <BNLowLevelILOperation N> size_t GetTrueTarget() const { return As<N>().GetTrueTarget(); }
 		template <BNLowLevelILOperation N> size_t GetFalseTarget() const { return As<N>().GetFalseTarget(); }
@@ -551,6 +553,7 @@ namespace BinaryNinja
 		SSARegister GetLowSSARegister() const;
 		int64_t GetConstant() const;
 		int64_t GetVector() const;
+		size_t GetStackAdjustment() const;
 		size_t GetTarget() const;
 		size_t GetTrueTarget() const;
 		size_t GetFalseTarget() const;
@@ -774,6 +777,11 @@ namespace BinaryNinja
 	{
 		LowLevelILInstruction GetDestExpr() const { return GetRawOperandAsExpr(0); }
 	};
+	template <> struct LowLevelILInstructionAccessor<LLIL_CALL_STACK_ADJUST>: public LowLevelILInstructionBase
+	{
+		LowLevelILInstruction GetDestExpr() const { return GetRawOperandAsExpr(0); }
+		size_t GetStackAdjustment() const { return (size_t)GetRawOperandAsInteger(1); }
+	};
 	template <> struct LowLevelILInstructionAccessor<LLIL_RET>: public LowLevelILInstructionBase
 	{
 		LowLevelILInstruction GetDestExpr() const { return GetRawOperandAsExpr(0); }
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index e50f80c6..75a3f1ad 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -161,6 +161,7 @@ class LowLevelILInstruction(object):
 		LowLevelILOperation.LLIL_JUMP: [("dest", "expr")],
 		LowLevelILOperation.LLIL_JUMP_TO: [("dest", "expr"), ("targets", "int_list")],
 		LowLevelILOperation.LLIL_CALL: [("dest", "expr")],
+		LowLevelILOperation.LLIL_CALL_STACK_ADJUST: [("dest", "expr"), ("stack_adjustment", "int")],
 		LowLevelILOperation.LLIL_RET: [("dest", "expr")],
 		LowLevelILOperation.LLIL_NORET: [],
 		LowLevelILOperation.LLIL_IF: [("condition", "expr"), ("true", "int"), ("false", "int")],
@@ -1262,6 +1263,18 @@ class LowLevelILFunction(object):
 		"""
 		return self.expr(LowLevelILOperation.LLIL_CALL, dest.index)
 
+	def call_stack_adjust(self, dest, stack_adjust):
+		"""
+		``call_stack_adjust`` returns an expression which first pushes the address of the next instruction onto the stack
+		then jumps (branches) to the expression ``dest``. After the function exits, ``stack_adjust`` is added to the
+		stack pointer register.
+
+		:param LowLevelILExpr dest: the expression to call
+		:return: The expression ``call(dest), stack += stack_adjust``
+		:rtype: LowLevelILExpr
+		"""
+		return self.expr(LowLevelILOperation.LLIL_CALL_STACK_ADJUST, dest.index, stack_adjust)
+
 	def ret(self, dest):
 		"""
 		``ret`` returns an expression which jumps (branches) to the expression ``dest``. ``ret`` is a special alias for
diff --git a/python/types.py b/python/types.py
index bf4a7b74..2557db2c 100644
--- a/python/types.py
+++ b/python/types.py
@@ -362,6 +362,12 @@ class Type(object):
 		"""Offset into structure (read-only)"""
 		return core.BNGetTypeOffset(self.handle)
 
+	@property
+	def stack_adjustment(self):
+		"""Stack adjustment for function (read-only)"""
+		result = core.BNGetTypeStackAdjustment(self.handle)
+		return SizeWithConfidence(result.value, confidence = result.confidence)
+
 	def __str__(self):
 		platform = None
 		if self.platform is not None:
@@ -551,7 +557,7 @@ class Type(object):
 		return Type(core.BNCreateArrayType(type_conf, count))
 
 	@classmethod
-	def function(self, ret, params, calling_convention=None, variable_arguments=None):
+	def function(self, ret, params, calling_convention=None, variable_arguments=None, stack_adjust=None):
 		"""
 		``function`` class method for creating an function Type.
 
@@ -605,8 +611,17 @@ class Type(object):
 		vararg_conf.value = variable_arguments.value
 		vararg_conf.confidence = variable_arguments.confidence
 
+		if stack_adjust is None:
+			stack_adjust = SizeWithConfidence(0, confidence = 0)
+		elif not isinstance(stack_adjust, SizeWithConfidence):
+			stack_adjust = SizeWithConfidence(stack_adjust)
+
+		stack_adjust_conf = core.BNSizeWithConfidence()
+		stack_adjust_conf.value = stack_adjust.value
+		stack_adjust_conf.confidence = stack_adjust.confidence
+
 		return Type(core.BNCreateFunctionType(ret_conf, conv_conf, param_buf, len(params),
-			vararg_conf))
+			vararg_conf, stack_adjust_conf))
 
 	@classmethod
 	def generate_auto_type_id(self, source, name):
diff --git a/type.cpp b/type.cpp
index aa891557..40a69ae2 100644
--- a/type.cpp
+++ b/type.cpp
@@ -425,6 +425,13 @@ uint64_t Type::GetOffset() const
 }
 
 
+Confidence<size_t> Type::GetStackAdjustment() const
+{
+	BNSizeWithConfidence result = BNGetTypeStackAdjustment(m_object);
+	return Confidence<size_t>(result.value, result.confidence);
+}
+
+
 string Type::GetString(Platform* platform) const
 {
 	char* str = BNGetTypeString(m_object, platform ? platform->GetObject() : nullptr);
@@ -663,7 +670,8 @@ Ref<Type> Type::ArrayType(const Confidence<Ref<Type>>& type, uint64_t elem)
 
 Ref<Type> Type::FunctionType(const Confidence<Ref<Type>>& returnValue,
 	const Confidence<Ref<CallingConvention>>& callingConvention,
-	const std::vector<FunctionParameter>& params, const Confidence<bool>& varArg)
+	const std::vector<FunctionParameter>& params, const Confidence<bool>& varArg,
+	const Confidence<size_t>& stackAdjust)
 {
 	BNTypeWithConfidence returnValueConf;
 	returnValueConf.type = returnValue->GetObject();
@@ -689,8 +697,12 @@ Ref<Type> Type::FunctionType(const Confidence<Ref<Type>>& returnValue,
 	varArgConf.value = varArg.GetValue();
 	varArgConf.confidence = varArg.GetConfidence();
 
+	BNSizeWithConfidence stackAdjustConf;
+	stackAdjustConf.value = stackAdjust.GetValue();
+	stackAdjustConf.confidence = stackAdjust.GetConfidence();
+
 	Type* type = new Type(BNCreateFunctionType(&returnValueConf, &callingConventionConf,
-		paramArray, params.size(), &varArgConf));
+		paramArray, params.size(), &varArgConf, &stackAdjustConf));
 	delete[] paramArray;
 	return type;
 }