expose overriding call types at individual call sites

5 files changed, 54 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index b409a0e6..469850fa 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -3048,17 +3048,20 @@ __attribute__ ((format (printf, 1, 2)))
 		std::vector<IndirectBranchInfo> GetIndirectBranches();
 		std::vector<IndirectBranchInfo> GetIndirectBranchesAt(Architecture* arch, uint64_t addr);
 
+		void SetAutoCallTypeAdjustment(Architecture* arch, uint64_t addr, const Confidence<Ref<Type>>& adjust);
 		void SetAutoCallStackAdjustment(Architecture* arch, uint64_t addr, const Confidence<int64_t>& adjust);
 		void SetAutoCallRegisterStackAdjustment(Architecture* arch, uint64_t addr,
 			const std::map<uint32_t, Confidence<int32_t>>& adjust);
 		void SetAutoCallRegisterStackAdjustment(Architecture* arch, uint64_t addr, uint32_t regStack,
 			const Confidence<int32_t>& adjust);
+		void SetUserCallTypeAdjustment(Architecture* arch, uint64_t addr, const Confidence<Ref<Type>>& adjust);
 		void SetUserCallStackAdjustment(Architecture* arch, uint64_t addr, const Confidence<int64_t>& adjust);
 		void SetUserCallRegisterStackAdjustment(Architecture* arch, uint64_t addr,
 			const std::map<uint32_t, Confidence<int32_t>>& adjust);
 		void SetUserCallRegisterStackAdjustment(Architecture* arch, uint64_t addr, uint32_t regStack,
 			const Confidence<int32_t>& adjust);
 
+		Confidence<Ref<Type>> GetCallTypeAdjustment(Architecture* arch, uint64_t addr);
 		Confidence<int64_t> GetCallStackAdjustment(Architecture* arch, uint64_t addr);
 		std::map<uint32_t, Confidence<int32_t>> GetCallRegisterStackAdjustment(Architecture* arch, uint64_t addr);
 		Confidence<int32_t> GetCallRegisterStackAdjustment(Architecture* arch, uint64_t addr, uint32_t regStack);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 81ecc8af..2d531c33 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2861,6 +2861,10 @@ __attribute__ ((format (printf, 1, 2)))
 	                                                                 uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeIndirectBranchList(BNIndirectBranchInfo* branches);
 
+	BINARYNINJACOREAPI void BNSetAutoCallTypeAdjustment(BNFunction* func, BNArchitecture* arch, uint64_t addr, 
+		BNTypeWithConfidence* type);
+	BINARYNINJACOREAPI void BNSetUserCallTypeAdjustment(BNFunction* func, BNArchitecture* arch, uint64_t addr,
+		BNTypeWithConfidence* type);
 	BINARYNINJACOREAPI void BNSetAutoCallStackAdjustment(BNFunction* func, BNArchitecture* arch, uint64_t addr,
 		int64_t adjust, uint8_t confidence);
 	BINARYNINJACOREAPI void BNSetUserCallStackAdjustment(BNFunction* func, BNArchitecture* arch, uint64_t addr,
@@ -2874,6 +2878,7 @@ __attribute__ ((format (printf, 1, 2)))
 	BINARYNINJACOREAPI void BNSetUserCallRegisterStackAdjustmentForRegisterStack(BNFunction* func,
 		BNArchitecture* arch, uint64_t addr, uint32_t regStack, int32_t adjust, uint8_t confidence);
 
+	BINARYNINJACOREAPI BNTypeWithConfidence BNGetCallTypeAdjustment(BNFunction* func, BNArchitecture* arch, uint64_t addr);
 	BINARYNINJACOREAPI BNOffsetWithConfidence BNGetCallStackAdjustment(BNFunction* func, BNArchitecture* arch, uint64_t addr);
 	BINARYNINJACOREAPI BNRegisterStackAdjustment* BNGetCallRegisterStackAdjustment(BNFunction* func,
 		BNArchitecture* arch, uint64_t addr, size_t* count);
diff --git a/function.cpp b/function.cpp
index 917b1556..c218a78e 100644
--- a/function.cpp
+++ b/function.cpp
@@ -1069,6 +1069,15 @@ vector<IndirectBranchInfo> Function::GetIndirectBranchesAt(Architecture* arch, u
 }
 
 
+void Function::SetAutoCallTypeAdjustment(Architecture* arch, uint64_t addr, const Confidence<Ref<Type>>& adjust)
+{
+	BNTypeWithConfidence apiObject;
+	apiObject.type = adjust ? adjust->GetObject() : nullptr;
+	apiObject.confidence = adjust.GetConfidence();
+	BNSetAutoCallTypeAdjustment(m_object, arch->GetObject(), addr, adjust ? &apiObject : nullptr);
+}
+
+
 void Function::SetAutoCallStackAdjustment(Architecture* arch, uint64_t addr, const Confidence<int64_t>& adjust)
 {
 	BNSetAutoCallStackAdjustment(m_object, arch->GetObject(), addr, adjust.GetValue(), adjust.GetConfidence());
@@ -1100,6 +1109,15 @@ void Function::SetAutoCallRegisterStackAdjustment(Architecture* arch, uint64_t a
 }
 
 
+void Function::SetUserCallTypeAdjustment(Architecture* arch, uint64_t addr, const Confidence<Ref<Type>>& adjust)
+{
+	BNTypeWithConfidence apiObject;
+	apiObject.type = adjust ? adjust->GetObject() : nullptr;
+	apiObject.confidence = adjust.GetConfidence();
+	BNSetUserCallTypeAdjustment(m_object, arch->GetObject(), addr, adjust ? &apiObject : nullptr);
+}
+
+
 void Function::SetUserCallStackAdjustment(Architecture* arch, uint64_t addr, const Confidence<int64_t>& adjust)
 {
 	BNSetUserCallStackAdjustment(m_object, arch->GetObject(), addr, adjust.GetValue(), adjust.GetConfidence());
@@ -1131,6 +1149,13 @@ void Function::SetUserCallRegisterStackAdjustment(Architecture* arch, uint64_t a
 }
 
 
+Confidence<Ref<Type>> Function::GetCallTypeAdjustment(Architecture* arch, uint64_t addr)
+{
+	BNTypeWithConfidence result = BNGetCallTypeAdjustment(m_object, arch->GetObject(), addr);
+	return Confidence<Ref<Type>>(result.type ? new Type(result.type) : nullptr, result.confidence);
+}
+
+
 Confidence<int64_t> Function::GetCallStackAdjustment(Architecture* arch, uint64_t addr)
 {
 	BNOffsetWithConfidence result = BNGetCallStackAdjustment(m_object, arch->GetObject(), addr);
diff --git a/python/function.py b/python/function.py
index 4707aa45..4ecece2a 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2225,6 +2225,17 @@ class Function(object):
 		core.BNSetAutoCallRegisterStackAdjustmentForRegisterStack(self.handle, arch.handle, addr, reg_stack,
 			adjust.value, adjust.confidence)
 
+	def set_call_type_adjustment(self, addr, adjust_type, arch=None):
+		if arch is None:
+			arch = self.arch
+		if adjust_type is None:
+			tc = None
+		else:
+			tc = core.BNTypeWithConfidence()
+			tc.type = adjust_type.handle
+			tc.confidence = adjust_type.confidence
+		core.BNSetUserCallTypeAdjustment(self.handle, arch.handle, addr, tc)
+
 	def set_call_stack_adjustment(self, addr, adjust, arch=None):
 		if arch is None:
 			arch = self.arch
@@ -2256,6 +2267,15 @@ class Function(object):
 		core.BNSetUserCallRegisterStackAdjustmentForRegisterStack(self.handle, arch.handle, addr, reg_stack,
 			adjust.value, adjust.confidence)
 
+	def get_call_type_adjustment(self, addr, arch=None):
+		if arch is None:
+			arch = self.arch
+		result = core.BNGetCallTypeAdjustment(self.handle, arch.handle, addr)
+		if result.type:
+			platform = self.platform
+			return types.Type(result.type, platform = platform, confidence = result.confidence)
+		return None
+
 	def get_call_stack_adjustment(self, addr, arch=None):
 		if arch is None:
 			arch = self.arch
diff --git a/ui/flowgraphwidget.h b/ui/flowgraphwidget.h
index 0d1c6355..dc358871 100644
--- a/ui/flowgraphwidget.h
+++ b/ui/flowgraphwidget.h
@@ -319,6 +319,7 @@ private Q_SLOTS:
 	void reanalyze();
 
 	void setStackAdjustment();
+	void setCallTypeAdjustment();
 
 	void editInstruction();
 	void instrEditDoneEvent();