Demangler plugin API

Closes #467

5 files changed, 347 insertions, 3 deletions

diff --git a/rust/Cargo.lock b/rust/Cargo.lock
index 3b8c11e2..a3c62e33 100644
--- a/rust/Cargo.lock
+++ b/rust/Cargo.lock
@@ -1040,6 +1040,14 @@ dependencies = [
 ]
 
 [[package]]
+name = "test_demangler"
+version = "0.1.0"
+dependencies = [
+ "binaryninja",
+ "log",
+]
+
+[[package]]
 name = "thiserror"
 version = "1.0.58"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 789a4224..7384893e 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -36,7 +36,8 @@ members = [
     "examples/hlil_lifter",
     "examples/pdb-ng",
     "examples/pdb-ng/demo",
-    "examples/template"
+    "examples/template",
+    "examples/test_demangler",
 ]
 
 [profile.release]
diff --git a/rust/examples/test_demangler/Cargo.toml b/rust/examples/test_demangler/Cargo.toml
new file mode 100644
index 00000000..eb4daa34
--- /dev/null
+++ b/rust/examples/test_demangler/Cargo.toml
@@ -0,0 +1,17 @@
+[package]
+name = "test_demangler"
+version = "0.1.0"
+edition = "2021"
+
+# Uncomment this if you're writing a plugin (plugins are shared objects loaded by the core):
+[lib]
+crate-type = ["cdylib"]
+
+# You can point at the BinaryNinja dependency in one of two ways, via path:
+[dependencies]
+binaryninja = {path="../../"}
+log = "0.4.21"
+
+# Or directly at the git repo:
+# [dependencies]
+# binaryninja = {git = "https://github.com/Vector35/binaryninja-api.git", branch = "dev"}
diff --git a/rust/examples/test_demangler/src/lib.rs b/rust/examples/test_demangler/src/lib.rs
new file mode 100644
index 00000000..eaa38dad
--- /dev/null
+++ b/rust/examples/test_demangler/src/lib.rs
@@ -0,0 +1,64 @@
+use log::{info, LevelFilter};
+use binaryninja::architecture::CoreArchitecture;
+use binaryninja::binaryview::BinaryView;
+use binaryninja::{command, logger};
+use binaryninja::command::Command;
+use binaryninja::demangle::{Demangler, CustomDemangler};
+use binaryninja::rc::Ref;
+use binaryninja::types::{QualifiedName, Type};
+
+struct TestDemangler;
+
+impl CustomDemangler for TestDemangler {
+    fn is_mangled_string(&self, name: &str) -> bool {
+        name == "test_name" || name == "test_name2"
+    }
+
+    fn demangle(&self, _arch: &CoreArchitecture, name: &str, _view: Option<Ref<BinaryView>>) -> Result<(Option<Ref<Type>>, QualifiedName), ()> {
+        match name {
+            "test_name" => Ok((Some(Type::bool()), QualifiedName::from(vec!["test_name"]))),
+            "test_name2" => Ok((None, QualifiedName::from(vec!["test_name2", "aaa"]))),
+            _ => Err(()),
+        }
+
+    }
+}
+
+struct DemangleCommand;
+
+impl Command for DemangleCommand {
+    fn action(&self, view: &BinaryView) {
+        for d in Demangler::list().iter() {
+            info!("{}", d.name());
+
+            info!("{}", d.is_mangled_string("__ZN1AC2Ei"));
+            info!("{:?}", d.demangle(
+                &CoreArchitecture::by_name("x86_64").expect("x86 exists"),
+                "__ZN1AC2Ei",
+                Some(view)
+            ));
+            info!("{:?}", d.demangle(
+                &CoreArchitecture::by_name("x86_64").expect("x86 exists"),
+                "test_name",
+                None
+            ));
+            info!("{:?}", d.demangle(
+                &CoreArchitecture::by_name("x86_64").expect("x86 exists"),
+                "test_name2",
+                None
+            ));
+        }
+    }
+
+    fn valid(&self, _view: &BinaryView) -> bool {
+        true
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn CorePluginInit() -> bool {
+    let _ = logger::init(LevelFilter::Info);
+    Demangler::register("Test", TestDemangler {});
+    command::register("Demangle Test", "Test", DemangleCommand {});
+    true
+}
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)>;
+}