Added some stack layout and reference APIs

5 files changed, 257 insertions, 1 deletions

diff --git a/rust/src/binaryview.rs b/rust/src/binaryview.rs
index 9e56f855..05f0d06a 100644
--- a/rust/src/binaryview.rs
+++ b/rust/src/binaryview.rs
@@ -22,6 +22,7 @@ use binaryninjacore_sys::*;
 pub use binaryninjacore_sys::BNModificationStatus as ModificationStatus;
 
 use std::ops;
+use std::ops::Range;
 use std::os::raw::c_char;
 use std::ptr;
 use std::result;
@@ -48,6 +49,7 @@ use crate::types::{DataVariable, NamedTypeReference, QualifiedName, QualifiedNam
 use crate::Endianness;
 
 use crate::rc::*;
+use crate::references::{CodeReference, DataReference};
 use crate::string::*;
 
 // TODO : general reorg of modules related to bv
@@ -1003,6 +1005,98 @@ pub trait BinaryViewExt: BinaryViewBase {
             )
         };
     }
+
+    /// Retrieves a list of [CodeReference]s pointing to a given address.
+    fn get_code_refs(&self, addr: u64) -> Array<CodeReference> {
+        unsafe {
+            let mut count = 0;
+            let handle = BNGetCodeReferences(self.as_ref().handle, addr, &mut count);
+            Array::new(handle, count, ())
+        }
+    }
+
+    /// Retrieves a list of [CodeReference]s pointing into a given [Range].
+    fn get_code_refs_in_range(&self, range: Range<u64>) -> Array<CodeReference> {
+        unsafe {
+            let mut count = 0;
+            let handle = BNGetCodeReferencesInRange(
+                self.as_ref().handle,
+                range.start,
+                range.end - range.start,
+                &mut count,
+            );
+            Array::new(handle, count, ())
+        }
+    }
+
+    /// Retrieves a list of [DataReference]s pointing to a given address.
+    fn get_data_refs(&self, addr: u64) -> Array<DataReference> {
+        unsafe {
+            let mut count = 0;
+            let handle = BNGetDataReferences(self.as_ref().handle, addr, &mut count);
+            Array::new(handle, count, ())
+        }
+    }
+
+    /// Retrieves a list of [DataReference]s originating from a given address.
+    fn get_data_refs_from(&self, addr: u64) -> Array<DataReference> {
+        unsafe {
+            let mut count = 0;
+            let handle = BNGetDataReferencesFrom(self.as_ref().handle, addr, &mut count);
+            Array::new(handle, count, ())
+        }
+    }
+
+    /// Retrieves a list of [DataReference]s pointing into a given [Range].
+    fn get_data_refs_in_range(&self, range: Range<u64>) -> Array<DataReference> {
+        unsafe {
+            let mut count = 0;
+            let handle = BNGetDataReferencesInRange(
+                self.as_ref().handle,
+                range.start,
+                range.end - range.start,
+                &mut count,
+            );
+            Array::new(handle, count, ())
+        }
+    }
+
+    /// Retrieves a list of [CodeReference]s for locations in code that use a given named type.
+    ///
+    /// TODO: It might be cleaner if this used an already allocated type from the core and
+    /// used its name instead of this slightly-gross [QualifiedName] hack. Since the returned
+    /// object doesn't have any [QualifiedName], I'm assuming the core does not alias
+    /// the [QualifiedName] we pass to it and it is safe to destroy it on [Drop], as in this function.
+    fn get_code_refs_for_type<B: BnStrCompatible>(&self, name: B) -> Array<CodeReference> {
+        unsafe {
+            let mut count = 0;
+            let q_name = &mut QualifiedName::from(name).0;
+            let handle = BNGetCodeReferencesForType(
+                self.as_ref().handle,
+                q_name as *mut BNQualifiedName,
+                &mut count,
+            );
+            Array::new(handle, count, ())
+        }
+    }
+    /// Retrieves a list of [DataReference]s instances of a given named type in data.
+    ///
+    /// TODO: It might be cleaner if this used an already allocated type from the core and
+    /// used its name instead of this slightly-gross [QualifiedName] hack. Since the returned
+    /// object doesn't have any [QualifiedName], I'm assuming the core does not alias
+    /// the [QualifiedName] we pass to it and it is safe to destroy it on [Drop], as in this function.
+    fn get_data_refs_for_type<B: BnStrCompatible>(&self, name: B) -> Array<DataReference> {
+        unsafe {
+            let mut count = 0;
+            let q_name = &mut QualifiedName::from(name).0;
+            let handle = BNGetDataReferencesForType(
+                self.as_ref().handle,
+                q_name as *mut BNQualifiedName,
+                &mut count,
+            );
+            Array::new(handle, count, ())
+        }
+    }
 }
 
 impl<T: BinaryViewBase> BinaryViewExt for T {}
diff --git a/rust/src/function.rs b/rust/src/function.rs
index 426f936b..ab504bb7 100644
--- a/rust/src/function.rs
+++ b/rust/src/function.rs
@@ -23,7 +23,7 @@ use crate::{
     llil,
     platform::Platform,
     symbol::Symbol,
-    types::{Conf, Type},
+    types::{Conf, NamedTypedVariable, Type},
 };
 
 use std::{fmt, mem};
@@ -258,6 +258,14 @@ impl Function {
             BNSetFunctionUserType(self.handle, t.handle);
         }
     }
+
+    pub fn stack_layout(&self) -> Array<NamedTypedVariable> {
+        let mut count = 0;
+        unsafe {
+            let variables = BNGetStackLayout(self.handle, &mut count);
+            Array::new(variables, count, ())
+        }
+    }
 }
 
 impl fmt::Debug for Function {
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index aa4f7efb..94e6ea40 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -155,6 +155,7 @@ pub mod logger;
 pub mod metadata;
 pub mod platform;
 pub mod rc;
+pub mod references;
 pub mod section;
 pub mod segment;
 pub mod settings;
diff --git a/rust/src/references.rs b/rust/src/references.rs
new file mode 100644
index 00000000..e3a73013
--- /dev/null
+++ b/rust/src/references.rs
@@ -0,0 +1,94 @@
+use crate::architecture::CoreArchitecture;
+use crate::function::Function;
+use crate::rc::{CoreArrayProvider, CoreArrayWrapper, CoreOwnedArrayProvider, Ref};
+use binaryninjacore_sys::{BNFreeCodeReferences, BNFreeDataReferences, BNReferenceSource};
+use std::mem::ManuallyDrop;
+
+/// A struct representing a single code cross-reference.
+/// Taking a cue from [LinearDisassemblyLine], this struct uses [ManuallyDrop] to
+/// prevent destructors from being run on the [Function] object allocated by
+/// the core in [BNGetCodeReferences] (et al). The reference is cleaned up on [Drop] of
+/// the enclosing array object.
+#[derive(Debug)]
+pub struct CodeReference {
+    arch: CoreArchitecture,
+    func: ManuallyDrop<Ref<Function>>,
+    pub address: u64,
+}
+
+/// A struct representing a single data cross-reference.
+/// Data references have no associated metadata, so this object has only
+/// a single [address] attribute.
+pub struct DataReference {
+    pub address: u64,
+}
+
+impl CodeReference {
+    pub(crate) unsafe fn new(handle: &BNReferenceSource) -> Self {
+        let func = ManuallyDrop::new(Function::from_raw((*handle).func));
+        let arch = CoreArchitecture::from_raw((*handle).arch);
+        let address = (*handle).addr;
+        Self {
+            func,
+            arch,
+            address,
+        }
+    }
+}
+
+impl<'a> CodeReference {
+    /// A handle to the referenced function bound by the [CodeReference] object's lifetime.
+    /// A user can call `.to_owned()` to promote this into its own ref-counted struct
+    /// and use it after the lifetime of the [CodeReference].
+    pub fn function(&'a self) -> &'a Function {
+        self.func.as_ref()
+    }
+
+    /// A handle to the [CodeReference]'s [CoreArchitecture]. This type is [Copy] so reference
+    /// shenanigans are not needed here.
+    pub fn architecture(&self) -> CoreArchitecture {
+        self.arch
+    }
+}
+
+// Code Reference Array<T> boilerplate
+
+impl CoreArrayProvider for CodeReference {
+    type Raw = BNReferenceSource;
+    type Context = ();
+}
+
+unsafe impl CoreOwnedArrayProvider for CodeReference {
+    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
+        BNFreeCodeReferences(raw, count)
+    }
+}
+
+unsafe impl<'a> CoreArrayWrapper<'a> for CodeReference {
+    type Wrapped = CodeReference;
+
+    unsafe fn wrap_raw(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped {
+        CodeReference::new(raw)
+    }
+}
+
+// Data Reference Array<T> boilerplate
+
+impl CoreArrayProvider for DataReference {
+    type Raw = u64;
+    type Context = ();
+}
+
+unsafe impl CoreOwnedArrayProvider for DataReference {
+    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
+        BNFreeDataReferences(raw)
+    }
+}
+
+unsafe impl<'a> CoreArrayWrapper<'a> for DataReference {
+    type Wrapped = DataReference;
+
+    unsafe fn wrap_raw(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped {
+        DataReference { address: *raw }
+    }
+}
diff --git a/rust/src/types.rs b/rust/src/types.rs
index 23a28631..174dcb14 100644
--- a/rust/src/types.rs
+++ b/rust/src/types.rs
@@ -38,6 +38,7 @@ use std::{
     hash::{Hash, Hasher},
     iter::{zip, IntoIterator},
     mem,
+    mem::ManuallyDrop,
     os::raw::c_char,
     ptr, result, slice,
     sync::Mutex,
@@ -1343,6 +1344,64 @@ impl Variable {
     }
 }
 
+///////////////
+// NamedVariable
+
+pub struct NamedTypedVariable {
+    var: BNVariable,
+    auto_defined: bool,
+    type_confidence: u8,
+    name: *mut std::os::raw::c_char,
+    ty: *mut BNType,
+}
+
+impl NamedTypedVariable {
+    pub fn name(&self) -> &str {
+        unsafe { BnStr::from_raw(self.name).as_str() }
+    }
+
+    pub fn var(&self) -> Variable {
+        unsafe { Variable::from_raw(self.var) }
+    }
+
+    pub fn auto_defined(&self) -> bool {
+        self.auto_defined
+    }
+
+    pub fn type_confidence(&self) -> u8 {
+        self.type_confidence
+    }
+
+    pub fn var_type(&self) -> Ref<Type> {
+        unsafe { Ref::new(Type::from_raw(self.ty)) }
+    }
+}
+
+impl CoreArrayProvider for NamedTypedVariable {
+    type Raw = BNVariableNameAndType;
+    type Context = ();
+}
+
+unsafe impl CoreOwnedArrayProvider for NamedTypedVariable {
+    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
+        BNFreeVariableNameAndTypeList(raw, count)
+    }
+}
+
+unsafe impl<'a> CoreArrayWrapper<'a> for NamedTypedVariable {
+    type Wrapped = ManuallyDrop<NamedTypedVariable>;
+
+    unsafe fn wrap_raw(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped {
+        ManuallyDrop::new(NamedTypedVariable {
+            var: raw.var,
+            ty: raw.type_,
+            name: raw.name,
+            auto_defined: raw.autoDefined,
+            type_confidence: raw.typeConfidence,
+        })
+    }
+}
+
 ////////////////////////
 // EnumerationBuilder