[Rust API] Helpers and Structure::members/insert_member

2 files changed, 205 insertions, 10 deletions

diff --git a/rust/src/string.rs b/rust/src/string.rs
index c4419b5e..e6de576b 100644
--- a/rust/src/string.rs
+++ b/rust/src/string.rs
@@ -22,6 +22,7 @@ use std::ops::Deref;
 use std::os::raw;
 
 use crate::rc::*;
+use crate::types::QualifiedName;
 
 pub(crate) fn raw_to_string(ptr: *const raw::c_char) -> Option<String> {
     if ptr.is_null() {
@@ -137,6 +138,17 @@ impl Drop for BnString {
     }
 }
 
+impl Clone for BnString {
+    fn clone(&self) -> Self {
+        use binaryninjacore_sys::BNAllocString;
+        unsafe {
+            Self {
+                raw: BNAllocString(self.raw),
+            }
+        }
+    }
+}
+
 impl Deref for BnString {
     type Target = BnStr;
 
@@ -157,6 +169,12 @@ impl fmt::Display for BnString {
     }
 }
 
+impl fmt::Debug for BnString {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.as_cstr().to_string_lossy())
+    }
+}
+
 impl CoreArrayProvider for BnString {
     type Raw = *mut raw::c_char;
     type Context = ();
@@ -239,3 +257,11 @@ unsafe impl<'a> BnStrCompatible for &'a Cow<'a, str> {
         self.as_ref().as_bytes()
     }
 }
+
+unsafe impl BnStrCompatible for &QualifiedName {
+    type Result = Vec<u8>;
+
+    fn as_bytes_with_nul(self) -> Self::Result {
+        self.string().as_bytes_with_nul()
+    }
+}
diff --git a/rust/src/types.rs b/rust/src/types.rs
index c1f24c59..9e7d20f2 100644
--- a/rust/src/types.rs
+++ b/rust/src/types.rs
@@ -17,7 +17,12 @@
 // TODO : Test the get_enumeration and get_structure methods
 
 use binaryninjacore_sys::*;
+use std::borrow::Cow;
+use std::ffi::CStr;
+use std::fmt::{Debug, Display, Formatter};
+use std::hash::{Hash, Hasher};
 use std::iter::IntoIterator;
+use std::os::raw::c_char;
 use std::{fmt, mem, ptr, result, slice};
 
 use crate::architecture::{Architecture, CoreArchitecture};
@@ -49,6 +54,38 @@ impl<T> Conf<T> {
             confidence,
         }
     }
+
+    pub fn map<U, F>(self, f: F) -> Conf<U>
+    where
+        F: FnOnce(T) -> U,
+    {
+        Conf::new(f(self.contents), self.confidence)
+    }
+
+    pub fn as_ref<U>(&self) -> Conf<&U>
+    where
+        T: AsRef<U>,
+    {
+        Conf::new(self.contents.as_ref(), self.confidence)
+    }
+}
+
+impl<T: Debug> Debug for Conf<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?} ({} confidence)", self.contents, self.confidence)
+    }
+}
+
+impl<T: Display> Display for Conf<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{} ({} confidence)", self.contents, self.confidence)
+    }
+}
+
+impl<'a, T> From<&'a Conf<T>> for Conf<&'a T> {
+    fn from(c: &'a Conf<T>) -> Self {
+        Conf::new(&c.contents, c.confidence)
+    }
 }
 
 #[inline]
@@ -1430,6 +1467,23 @@ impl StructureBuilder {
         self
     }
 
+    pub fn insert_member<'a, 'b>(
+        &'a mut self,
+        member: &'b StructureMember,
+        overwrite_existing: bool,
+    ) -> &'a mut Self {
+        let ty = member.ty.clone();
+        self.insert(
+            ty.as_ref(),
+            member.name.clone(),
+            member.offset,
+            overwrite_existing,
+            member.access,
+            member.scope,
+        );
+        self
+    }
+
     pub fn insert<'a, 'b, S: BnStrCompatible, T: Into<Conf<&'b Type>>>(
         &'a mut self,
         t: T,
@@ -1545,6 +1599,26 @@ impl Structure {
         unsafe { BNGetStructureType(self.handle) }
     }
 
+    pub fn members(&self) -> Result<Vec<StructureMember>> {
+        unsafe {
+            let mut count: usize = mem::zeroed();
+            let members_raw: *mut BNStructureMember =
+                BNGetStructureMembers(self.handle, &mut count);
+            if members_raw.is_null() {
+                return Err(());
+            }
+            let members = slice::from_raw_parts(members_raw, count);
+
+            let result = (0..count)
+                .map(|i| StructureMember::from_raw(members[i]))
+                .collect();
+
+            BNFreeStructureMemberList(members_raw, count);
+
+            Ok(result)
+        }
+    }
+
     // TODO : The other methods in the python version (alignment, packed, type, members, remove, replace, etc)
 }
 
@@ -1554,6 +1628,18 @@ impl From<&StructureBuilder> for Ref<Structure> {
     }
 }
 
+impl Debug for Structure {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "Structure {{")?;
+        if let Ok(members) = self.members() {
+            for member in members {
+                write!(f, " {:?}", member)?;
+            }
+        }
+        write!(f, "}}")
+    }
+}
+
 unsafe impl RefCountable for Structure {
     unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
         Ref::new(Self::from_raw(BNNewStructureReference(handle.handle)))
@@ -1572,15 +1658,45 @@ impl ToOwned for Structure {
     }
 }
 
-// TODO : Use
-// #[derive(PartialEq, Eq, Hash)]
-// pub struct StructureMember {
-//     pub(crate) handle: *mut BNStructureMember,
-// }
+#[derive(Debug, Clone)]
+pub struct StructureMember {
+    pub ty: Conf<Ref<Type>>,
+    pub name: BnString,
+    pub offset: u64,
+    pub access: MemberAccess,
+    pub scope: MemberScope,
+}
 
-// impl StructureMember {
-//     // pub fn new() -> Self {}
-// }
+impl StructureMember {
+    pub fn new<T: BnStrCompatible>(
+        ty: Conf<Ref<Type>>,
+        name: T,
+        offset: u64,
+        access: MemberAccess,
+        scope: MemberScope,
+    ) -> Self {
+        Self {
+            ty,
+            name: BnString::new(name),
+            offset,
+            access,
+            scope,
+        }
+    }
+
+    pub(crate) unsafe fn from_raw(handle: BNStructureMember) -> Self {
+        Self {
+            ty: Conf::new(
+                RefCountable::inc_ref(&Type::from_raw(handle.type_)),
+                handle.typeConfidence,
+            ),
+            name: BnString::new(BnStr::from_raw(handle.name)),
+            offset: handle.offset,
+            access: handle.access,
+            scope: handle.scope,
+        }
+    }
+}
 
 ////////////////////////
 // NamedTypeReference
@@ -1610,6 +1726,11 @@ impl NamedTypeReference {
             },
         }
     }
+
+    pub fn name(&self) -> QualifiedName {
+        let named_ref: BNQualifiedName = unsafe { BNGetTypeReferenceName(self.handle) };
+        QualifiedName(named_ref)
+    }
 }
 
 ///////////////////
@@ -1621,8 +1742,6 @@ pub struct QualifiedName(pub(crate) BNQualifiedName);
 impl QualifiedName {
     // TODO : I think this is bad
     pub fn string(&self) -> String {
-        use std::ffi::CStr;
-
         unsafe {
             slice::from_raw_parts(self.0.name, self.0.nameCount)
                 .iter()
@@ -1631,6 +1750,21 @@ impl QualifiedName {
                 .join("::")
         }
     }
+
+    pub fn join(&self) -> Cow<str> {
+        let join: *mut c_char = self.0.join;
+        unsafe { CStr::from_ptr(join).to_string_lossy() }
+    }
+
+    pub fn strings(&self) -> Vec<Cow<str>> {
+        let names: *mut *mut c_char = self.0.name;
+        unsafe {
+            slice::from_raw_parts(names, self.0.nameCount)
+                .iter()
+                .map(|name| CStr::from_ptr(*name).to_string_lossy())
+                .collect::<Vec<_>>()
+        }
+    }
 }
 
 impl<S: BnStrCompatible> From<S> for QualifiedName {
@@ -1667,6 +1801,41 @@ impl<S: BnStrCompatible> From<Vec<S>> for QualifiedName {
     }
 }
 
+impl Clone for QualifiedName {
+    fn clone(&self) -> Self {
+        let strings = self.strings();
+        let name = Self::from(strings.iter().collect::<Vec<&Cow<str>>>());
+        name
+    }
+}
+
+impl Hash for QualifiedName {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.join().hash(state);
+        self.strings().hash(state);
+    }
+}
+
+impl Debug for QualifiedName {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.string())
+    }
+}
+
+impl Display for QualifiedName {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.string())
+    }
+}
+
+impl PartialEq for QualifiedName {
+    fn eq(&self, other: &Self) -> bool {
+        self.strings() == other.strings()
+    }
+}
+
+impl Eq for QualifiedName {}
+
 impl Drop for QualifiedName {
     fn drop(&mut self) {
         unsafe {