Support adding user-defined xrefs

6 files changed, 83 insertions, 28 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 10c64642..fe763064 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1413,10 +1413,13 @@ namespace BinaryNinja
 
 		std::vector<ReferenceSource> GetCodeReferences(uint64_t addr);
 		std::vector<ReferenceSource> GetCodeReferences(uint64_t addr, uint64_t len);
+
 		std::vector<uint64_t> GetDataReferences(uint64_t addr);
 		std::vector<uint64_t> GetDataReferences(uint64_t addr, uint64_t len);
 		std::vector<uint64_t> GetDataReferencesFrom(uint64_t addr);
 		std::vector<uint64_t> GetDataReferencesFrom(uint64_t addr, uint64_t len);
+		void AddUserDataReference(uint64_t fromAddr, uint64_t toAddr);
+		void RemoveUserDataReference(uint64_t fromAddr, uint64_t toAddr);
 
 		Ref<Symbol> GetSymbolByAddress(uint64_t addr, const NameSpace& nameSpace=NameSpace());
 		Ref<Symbol> GetSymbolByRawName(const std::string& name, const NameSpace& nameSpace=NameSpace());
@@ -2588,8 +2591,8 @@ namespace BinaryNinja
 		void SetComment(const std::string& comment);
 		void SetCommentForAddress(uint64_t addr, const std::string& comment);
 
-		void SetUserXref(uint64_t addr, uint64_t target);
-		void RemoveUserXref(uint64_t addr, uint64_t target);
+		void AddUserCodeRef(Architecture* fromArch, uint64_t fromAddr, uint64_t toAddr);
+		void RemoveUserCodeRef(Architecture* fromArch, uint64_t fromAddr, uint64_t toAddr);
 
 		Ref<LowLevelILFunction> GetLowLevelIL() const;
 		size_t GetLowLevelILForInstruction(Architecture* arch, uint64_t addr);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 628f53ed..174d09e0 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2528,8 +2528,8 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetFunctionComment(BNFunction* func, const char* comment);
 	BINARYNINJACOREAPI void BNSetCommentForAddress(BNFunction* func, uint64_t addr, const char* comment);
 
-	BINARYNINJACOREAPI void BNSetUserXref(BNFunction* func, uint64_t addr, uint64_t target);
-	BINARYNINJACOREAPI void BNRemoveUserXref(BNFunction* func, uint64_t addr, uint64_t target);
+	BINARYNINJACOREAPI void BNAddUserCodeRef(BNFunction* func, BNArchitecture* fromArch, uint64_t fromAddr, uint64_t toAddr);
+	BINARYNINJACOREAPI void BNRemoveUserCodeRef(BNFunction* func, BNArchitecture* fromArch, uint64_t fromAddr, uint64_t toAddr);
 
 	BINARYNINJACOREAPI BNBasicBlock* BNNewBasicBlockReference(BNBasicBlock* block);
 	BINARYNINJACOREAPI void BNFreeBasicBlock(BNBasicBlock* block);
@@ -2695,10 +2695,13 @@ extern "C"
 	BINARYNINJACOREAPI BNReferenceSource* BNGetCodeReferencesInRange(BNBinaryView* view, uint64_t addr,
 	                                                                 uint64_t len, size_t* count);
 	BINARYNINJACOREAPI void BNFreeCodeReferences(BNReferenceSource* refs, size_t count);
+
 	BINARYNINJACOREAPI uint64_t* BNGetDataReferences(BNBinaryView* view, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI uint64_t* BNGetDataReferencesInRange(BNBinaryView* view, uint64_t addr, uint64_t len, size_t* count);
 	BINARYNINJACOREAPI uint64_t* BNGetDataReferencesFrom(BNBinaryView* view, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI uint64_t* BNGetDataReferencesFromInRange(BNBinaryView* view, uint64_t addr, uint64_t len, size_t* count);
+	BINARYNINJACOREAPI void BNAddUserDataReference(BNBinaryView* view, uint64_t fromAddr, uint64_t toAddr);
+	BINARYNINJACOREAPI void BNRemoveUserDataReference(BNBinaryView* view, uint64_t fromAddr, uint64_t toAddr);
 	BINARYNINJACOREAPI void BNFreeDataReferences(uint64_t* refs);
 
 	BINARYNINJACOREAPI void BNRegisterGlobalFunctionRecognizer(BNFunctionRecognizer* rec);
diff --git a/binaryview.cpp b/binaryview.cpp
index e0948a1c..547d1fc7 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -1482,6 +1482,18 @@ vector<uint64_t> BinaryView::GetDataReferencesFrom(uint64_t addr, uint64_t len)
 }
 
 
+void BinaryView::AddUserDataReference(uint64_t fromAddr, uint64_t toAddr)
+{
+	BNAddUserDataReference(m_object, fromAddr, toAddr);
+}
+
+
+void BinaryView::RemoveUserDataReference(uint64_t fromAddr, uint64_t toAddr)
+{
+	BNRemoveUserDataReference(m_object, fromAddr, toAddr);
+}
+
+
 Ref<Symbol> BinaryView::GetSymbolByAddress(uint64_t addr, const NameSpace& nameSpace)
 {
 	BNNameSpace ns = nameSpace.GetAPIObject();
diff --git a/function.cpp b/function.cpp
index 54c91541..2a3e333f 100644
--- a/function.cpp
+++ b/function.cpp
@@ -246,15 +246,15 @@ void Function::SetCommentForAddress(uint64_t addr, const string& comment)
 }
 
 
-void Function::SetUserXref(uint64_t addr, uint64_t target)
+void Function::AddUserCodeRef(Architecture* fromArch, uint64_t fromAddr, uint64_t toAddr)
 {
-	BNSetUserXref(m_object, addr, target);
+	BNAddUserCodeRef(m_object, fromArch->GetObject(), fromAddr, toAddr);
 }
 
 
-void Function::RemoveUserXref(uint64_t addr, uint64_t target)
+void Function::RemoveUserCodeRef(Architecture* fromArch, uint64_t fromAddr, uint64_t toAddr)
 {
-	BNRemoveUserXref(m_object, addr, target);
+	BNRemoveUserCodeRef(m_object, fromArch->GetObject(), fromAddr, toAddr);
 }
 
 
diff --git a/python/binaryview.py b/python/binaryview.py
index f3ee7733..bf55c047 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -2745,6 +2745,8 @@ class BinaryView(object):
 	def get_code_refs(self, addr, length=None):
 		"""
 		``get_code_refs`` returns a list of ReferenceSource objects (xrefs or cross-references) that point to the provided virtual address.
+		This function returns both autoanalysis ("auto") and user-specified ("user") xrefs. 
+		To add a user-specified reference, see :func:`~binaryninja.function.Function.add_user_code_ref`.
 
 		:param int addr: virtual address to query for references
 		:return: List of References for the given virtual address
@@ -2780,6 +2782,8 @@ class BinaryView(object):
 		"""
 		``get_data_refs`` returns a list of virtual addresses of data which references ``addr``. Optionally specifying
 		a length. When ``length`` is set ``get_data_refs`` returns the data which references in the range ``addr``-``addr``+``length``.
+		This function returns both autoanalysis ("auto") and user-specified ("user") xrefs. To add a user-specified
+		reference, see :func:`add_user_data_ref`.
 
 		:param int addr: virtual address to query for references
 		:param int length: optional length of query
@@ -2807,6 +2811,8 @@ class BinaryView(object):
 		"""
 		``get_data_refs_from`` returns a list of virtual addresses referenced by the address ``addr``. Optionally specifying
 		a length. When ``length`` is set ``get_data_refs_from`` returns the data referenced in the range ``addr``-``addr``+``length``.
+		This function returns both autoanalysis ("auto") and user-specified ("user") xrefs. To add a user-specified
+		reference, see :func:`add_user_data_ref`.
 
 		:param int addr: virtual address to query for references
 		:param int length: optional length of query
@@ -2831,6 +2837,31 @@ class BinaryView(object):
 		return result
 
 
+	def add_user_data_ref(self, from_addr, to_addr):
+		"""
+		``add_user_data_ref`` adds a user-specified data cross-reference (xref) from the address ``from_addr`` to the address ``to_addr``.
+		If the reference already exists, no action is performed. To remove the reference, use :func:`remove_user_data_ref`.
+
+		:param int from_addr: the reference's source virtual address.
+		:param int to_addr: the reference's destination virtual address.
+		:rtype: None
+		"""
+		core.BNAddUserDataReference(self.handle, from_addr, to_addr)
+
+
+	def remove_user_data_ref(self, from_addr, to_addr):
+		"""
+		``remove_user_data_ref`` removes a user-specified data cross-reference (xref) from the address ``from_addr`` to the address ``to_addr``.
+		This function will only remove user-specified references, not ones generated during autoanalysis.
+		If the reference does not exist, no action is performed.
+
+		:param int from_addr: the reference's source virtual address.
+		:param int to_addr: the reference's destination virtual address.
+		:rtype: None
+		"""
+		core.BNRemoveUserDataReference(self.handle, from_addr, to_addr)
+
+
 	def get_symbol_at(self, addr, namespace=None):
 		"""
 		``get_symbol_at`` returns the Symbol at the provided virtual address.
diff --git a/python/function.py b/python/function.py
index 003b5a50..4aa9cc39 100644
--- a/python/function.py
+++ b/python/function.py
@@ -1293,42 +1293,48 @@ class Function(object):
 		"""
 		core.BNSetCommentForAddress(self.handle, addr, comment)
 
-	def set_user_xref(self, addr, target):
+	def add_user_code_ref(self, from_addr, to_addr, from_arch=None):
 		"""
-		``set_user_xref`` places a user-defined cross-reference from the given address to
-		the specified target address.
-		If the given address is not contained within this function, no action is performed.
-		If there are multiple basic blocks (i.e. of different architectures) which contain
-		the specified source address, then the cross-reference is added for all matching
-		matching basic blocks. To remove the xref, use ``remove_user_xref``.
+		``add_user_code_ref`` places a user-defined cross-reference from the instruction at
+		the given address and architecture to the specified target address. If the specified
+		source instruction is not contained within this function, no action is performed.
+		To remove the reference, use :func:`remove_user_code_ref`.
 
-		:param addr int: virtual address within the current function of the xref's source.
-		:param target int: virtual address of the xref's destination.
+		:param from_addr int: virtual address of the source instruction
+		:param to_addr int: virtual address of the xref's destination.
+		:param from_arch Architecture: (optional) architecture of the source instruction
 		:rtype: None
 		:Example:
 
-			>>> current_function.set_user_xref(here, 0x400000)
+			>>> current_function.add_user_code_ref(here, 0x400000)
 
 		"""
-		core.BNSetUserXref(self.handle, addr, target)
 
-	def remove_user_xref(self, addr, target):
+		if from_arch is None:
+			from_arch = self.arch
+
+		core.BNAddUserCodeRef(self.handle, from_arch.handle, from_addr, to_addr)
+
+	def remove_user_code_ref(self, from_addr, to_addr, from_arch=None):
 		"""
-		``remove_user_xref`` reomves a user-defined cross-reference.
+		``remove_user_code_ref`` reomves a user-defined cross-reference.
 		If the given address is not contained within this function, or if there is no
-		user-defined cross-reference from the specified source address, no action is performed.
-		If there are multiple user-defined xrefs in the function which match the source and
-		target address, all matching xrefs are removed.
+		such user-defined cross-reference, no action is performed.
 
-		:param addr int: virtual address within the current function of the xref's source.
-		:param target int: virtual address of the xref's destination.
+		:param from_addr int: virtual address of the source instruction
+		:param to_addr int: virtual address of the xref's destination.
+		:param from_arch Architecture: (optional) architecture of the source instruction
 		:rtype: None
 		:Example:
 
-			>>> current_function.remove_user_xref(here, 0x400000)
+			>>> current_function.remove_user_code_ref(here, 0x400000)
 
 		"""
-		core.BNRemoveUserXref(self.handle, addr, target)
+
+		if from_arch is None:
+			from_arch = self.arch
+
+		core.BNRemoveUserCodeRef(self.handle, from_arch.handle, from_addr, to_addr)
 
 	def get_low_level_il_at(self, addr, arch=None):
 		"""