Add BinaryView.get_code_refs_from

4 files changed, 66 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index fe763064..8d34bef3 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1413,6 +1413,8 @@ namespace BinaryNinja
 
 		std::vector<ReferenceSource> GetCodeReferences(uint64_t addr);
 		std::vector<ReferenceSource> GetCodeReferences(uint64_t addr, uint64_t len);
+		std::vector<uint64_t> GetCodeReferencesFrom(ReferenceSource src);
+		std::vector<uint64_t> GetCodeReferencesFrom(ReferenceSource src, uint64_t len);
 
 		std::vector<uint64_t> GetDataReferences(uint64_t addr);
 		std::vector<uint64_t> GetDataReferences(uint64_t addr, uint64_t len);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 174d09e0..d7c43150 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2695,6 +2695,9 @@ 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* BNGetCodeReferencesFrom(BNBinaryView* view, BNReferenceSource* src, size_t* count);
+	BINARYNINJACOREAPI uint64_t* BNGetCodeReferencesFromInRange(BNBinaryView* view, BNReferenceSource* src, uint64_t len, size_t* count);
+	BINARYNINJACOREAPI void BNFreeCodeReferencesFrom(uint64_t* refs);
 
 	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);
diff --git a/binaryview.cpp b/binaryview.cpp
index 547d1fc7..c78e43cf 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -1442,6 +1442,38 @@ vector<ReferenceSource> BinaryView::GetCodeReferences(uint64_t addr, uint64_t le
 }
 
 
+vector<uint64_t> BinaryView::GetCodeReferencesFrom(ReferenceSource src)
+{
+	size_t count;
+	BNReferenceSource _src{ src.func->m_object, src.arch->m_object, src.addr };
+	uint64_t* refs = BNGetCodeReferencesFrom(m_object, &_src, &count);
+
+	vector<uint64_t> result;
+	result.reserve(count);
+	for (size_t i = 0; i < count; i++)
+		result[i] = refs[i];
+
+	BNFreeCodeReferencesFrom(refs);
+	return result;
+}
+
+
+vector<uint64_t> BinaryView::GetCodeReferencesFrom(ReferenceSource src, uint64_t len)
+{
+	size_t count;
+	BNReferenceSource _src{ src.func->m_object, src.arch->m_object, src.addr };
+	uint64_t* refs = BNGetCodeReferencesFromInRange(m_object, &_src, len, &count);
+
+	vector<uint64_t> result;
+	result.reserve(count);
+	for (size_t i = 0; i < count; i++)
+		result[i] = refs[i];
+
+	BNFreeCodeReferencesFrom(refs);
+	return result;
+}
+
+
 vector<uint64_t> BinaryView::GetDataReferences(uint64_t addr)
 {
 	size_t count;
diff --git a/python/binaryview.py b/python/binaryview.py
index 2f7d39be..ac82dbc3 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -2797,6 +2797,35 @@ class BinaryView(object):
 		core.BNFreeCodeReferences(refs, count.value)
 		return result
 
+	def get_code_refs_from(self, addr, func=None, arch=None, length=None):
+		"""
+		``get_code_refs_from`` returns a list of virtual addresses referenced by code in the function ``func``,
+		of the architecture ``arch``, and at the address ``addr``. If no function is specified, references from
+		all functions and containing the address will be returned. If no architecture is specified, the
+		architecture of the function will be used. 
+		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`.
+
+		todo: docs and make it not so difficult to call
+		"""
+
+		result = []
+		funcs = self.get_functions_containing(addr) if func is None else [func]
+		if not funcs:
+			return []
+		for src_func in funcs:
+			src_arch = src_func.arch if arch is None else arch
+			ref_src = core.BNReferenceSource(src_func.handle, src_arch.handle, addr)
+			count = ctypes.c_ulonglong(0)
+			if length is None:
+				refs = core.BNGetCodeReferencesFrom(self.handle, ref_src, count)
+			else:
+				refs = core.BNGetCodeReferencesFromInRange(self.handle, ref_src, length, count)
+			for i in range(0, count.value):
+				result.append(refs[i])
+			core.BNFreeCodeReferencesFrom(refs, count.value)
+		return result
+
 	def get_data_refs(self, addr, length=None):
 		"""
 		``get_data_refs`` returns a list of virtual addresses of data which references ``addr``. Optionally specifying