1 files changed, 256 insertions, 2 deletions
diff --git a/rust/src/demangle.rs b/rust/src/demangle.rs
index 1b940ff9..8fbc4316 100644
--- a/rust/src/demangle.rs
+++ b/rust/src/demangle.rs
@@ -17,15 +17,75 @@
 use binaryninjacore_sys::*;
 use std::os::raw::c_char;
 use std::{ffi::CStr, result};
+use std::ffi::c_void;
 
 use crate::architecture::CoreArchitecture;
-use crate::string::{BnStrCompatible, BnString};
-use crate::types::Type;
+use crate::binaryview::BinaryView;
+use crate::string::{BnStrCompatible, BnString, raw_to_string};
+use crate::types::{QualifiedName, Type};
 
 use crate::rc::*;
 
 pub type Result<R> = result::Result<R, ()>;
 
+pub fn demangle_generic<S: BnStrCompatible>(
+    arch: &CoreArchitecture,
+    mangled_name: S,
+    view: Option<&BinaryView>,
+    simplify: bool,
+) -> Result<(Option<Ref<Type>>, Vec<String>)> {
+    let mangled_name_bwn = mangled_name.into_bytes_with_nul();
+    let mangled_name_ptr = mangled_name_bwn.as_ref();
+    let mut out_type: *mut BNType = std::ptr::null_mut();
+    let mut out_name = BNQualifiedName {
+        name: std::ptr::null_mut(),
+        join: std::ptr::null_mut(),
+        nameCount: 0,
+    };
+    let view_ptr = match view {
+        Some(v) => v.handle,
+        None => std::ptr::null_mut(),
+    };
+    let res = unsafe {
+        BNDemangleGeneric(
+            arch.0,
+            mangled_name_ptr.as_ptr() as *const c_char,
+            &mut out_type,
+            &mut out_name,
+            view_ptr,
+            simplify,
+        )
+    };
+
+    if !res {
+        let cstr = match CStr::from_bytes_with_nul(mangled_name_ptr) {
+            Ok(cstr) => cstr,
+            Err(_) => {
+                log::error!("demangle_generic: failed to parse mangled name");
+                return Err(());
+            }
+        };
+        return Ok((None, vec![cstr.to_string_lossy().into_owned()]));
+    }
+
+    let out_type = match out_type.is_null() {
+        true => {
+            log::debug!("demangle_generic: out_type is NULL");
+            None
+        }
+        false => Some(unsafe { Type::ref_from_raw(out_type) }),
+    };
+
+    Ok((
+        out_type,
+        QualifiedName(out_name)
+            .strings()
+            .iter()
+            .map(|str| str.to_string())
+            .collect::<Vec<_>>()
+    ))
+}
+
 pub fn demangle_llvm<S: BnStrCompatible>(
     mangled_name: S,
     simplify: bool,
@@ -179,3 +239,197 @@ pub fn demangle_ms<S: BnStrCompatible>(
 
     Ok((out_type, names))
 }
+
+#[derive(PartialEq, Eq, Hash)]
+pub struct Demangler {
+    pub(crate) handle: *mut BNDemangler,
+}
+
+impl Demangler {
+    pub(crate) unsafe fn from_raw(handle: *mut BNDemangler) -> Self {
+        debug_assert!(!handle.is_null());
+
+        Self { handle }
+    }
+
+    pub fn is_mangled_string<S: BnStrCompatible>(&self, name: S) -> bool {
+        let bytes = name.into_bytes_with_nul();
+        unsafe { BNIsDemanglerMangledName(self.handle, bytes.as_ref().as_ptr() as *const _) }
+    }
+
+    pub fn demangle<S: BnStrCompatible>(
+        &self,
+        arch: &CoreArchitecture,
+        name: S,
+        view: Option<&BinaryView>,
+    ) -> Result<(Option<Ref<Type>>, QualifiedName)> {
+        let name_bytes = name.into_bytes_with_nul();
+
+        let mut out_type = std::ptr::null_mut();
+        let mut out_var_name = BNQualifiedName {
+            name: std::ptr::null_mut(),
+            join: std::ptr::null_mut(),
+            nameCount: 0,
+        };
+
+        let view_ptr = match view {
+            Some(v) => v.handle,
+            None => std::ptr::null_mut()
+        };
+
+        if !unsafe { BNDemanglerDemangle(self.handle, arch.0, name_bytes.as_ref().as_ptr() as *const _, &mut out_type, &mut out_var_name, view_ptr) } {
+            return Err(());
+        }
+
+        let var_type =
+            if out_type.is_null() {
+                None
+            } else {
+                Some(unsafe { Type::ref_from_raw(out_type) })
+            };
+        let var_name = QualifiedName(out_var_name);
+
+        Ok((var_type, var_name))
+    }
+
+    pub fn name(&self) -> BnString {
+        unsafe { BnString::from_raw(BNGetDemanglerName(self.handle)) }
+    }
+
+    pub fn from_name<S: BnStrCompatible>(name: S) -> Option<Self> {
+        let name_bytes = name.into_bytes_with_nul();
+        let demangler = unsafe { BNGetDemanglerByName(name_bytes.as_ref().as_ptr() as *const _) };
+        if demangler.is_null() {
+            None
+        } else {
+            Some(unsafe { Demangler::from_raw(demangler) })
+        }
+    }
+
+    pub fn list() -> Array<Self> {
+        let mut count: usize = 0;
+        let demanglers = unsafe { BNGetDemanglerList(&mut count) };
+        unsafe { Array::<Demangler>::new(demanglers, count, ()) }
+    }
+
+    pub fn register<S, C>(name: S, callbacks: C) -> Self
+    where
+        S: BnStrCompatible,
+        C: CustomDemangler,
+    {
+        extern "C" fn cb_is_mangled_string<C>(ctxt: *mut c_void, name: *const c_char) -> bool
+        where
+            C: CustomDemangler,
+        {
+            ffi_wrap!("CustomDemangler::cb_is_mangled_string", unsafe {
+                let cmd = &*(ctxt as *const C);
+                let name =
+                    if let Some(n) = raw_to_string(name) {
+                        n
+                    } else {
+                        return false;
+                    };
+                cmd.is_mangled_string(&name)
+            })
+        }
+        extern "C" fn cb_demangle<C>(ctxt: *mut c_void, arch: *mut BNArchitecture, name: *const c_char, out_type: *mut *mut BNType, out_var_name: *mut BNQualifiedName, view: *mut BNBinaryView) -> bool
+        where
+            C: CustomDemangler,
+        {
+            ffi_wrap!("CustomDemangler::cb_demangle", unsafe {
+                let cmd = &*(ctxt as *const C);
+                let arch = CoreArchitecture::from_raw(arch);
+                let name =
+                    if let Some(n) = raw_to_string(name) {
+                        n
+                    } else {
+                        return false;
+                    };
+                let view = if view.is_null() {
+                    None
+                } else {
+                    Some(BinaryView::from_raw(BNNewViewReference(view)))
+                };
+
+                match cmd.demangle(&arch, &name, view) {
+                    Ok((type_, name)) => {
+                        *out_type = match type_ {
+                            Some(t) => RefCountable::inc_ref(t.as_ref()).handle,
+                            None => std::ptr::null_mut()
+                        };
+                        // TODO: Need to have a better way for api-owned QNames
+                        (*out_var_name).nameCount = name.0.nameCount;
+                        (*out_var_name).join = BNAllocString(name.0.join);
+                        (*out_var_name).name = BNAllocStringList(name.0.name as *mut *const _, name.0.nameCount);
+                        true
+                    },
+                    Err(_) => {
+                        false
+                    }
+                }
+            })
+        }
+        extern "C" fn cb_free_var_name<C>(_ctxt: *mut c_void, name: *mut BNQualifiedName)
+        where
+            C: CustomDemangler,
+        {
+            ffi_wrap!("CustomDemangler::cb_free_var_name", unsafe {
+                BNFreeString((*name).join);
+                BNFreeStringList((*name).name, (*name).nameCount);
+            })
+        }
+
+        let name = name.into_bytes_with_nul();
+        let name_ptr = name.as_ref().as_ptr() as *mut _;
+        let ctxt = Box::into_raw(Box::new(callbacks));
+
+        let callbacks = BNDemanglerCallbacks {
+            context: ctxt as *mut c_void,
+            isMangledString: Some(cb_is_mangled_string::<C>),
+            demangle: Some(cb_demangle::<C>),
+            freeVarName: Some(cb_free_var_name::<C>),
+        };
+
+        unsafe {
+            Demangler::from_raw(BNRegisterDemangler(
+                name_ptr,
+                Box::leak(Box::new(callbacks)),
+            ))
+        }
+    }
+
+    pub fn promote(demangler: &Demangler) {
+        unsafe { BNPromoteDemangler(demangler.handle); }
+    }
+}
+
+unsafe impl Sync for Demangler {}
+
+unsafe impl Send for Demangler {}
+
+impl CoreArrayProvider for Demangler {
+    type Raw = *mut BNDemangler;
+    type Context = ();
+    type Wrapped<'a> = Demangler;
+}
+
+unsafe impl CoreArrayProviderInner for Demangler {
+    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
+        BNFreeDemanglerList(raw);
+    }
+
+    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
+        Demangler::from_raw(*raw)
+    }
+}
+
+pub trait CustomDemangler: 'static + Sync {
+    fn is_mangled_string(&self, name: &str) -> bool;
+
+    fn demangle(
+        &self,
+        arch: &CoreArchitecture,
+        name: &str,
+        view: Option<Ref<BinaryView>>,
+    ) -> Result<(Option<Ref<Type>>, QualifiedName)>;
+}