Support function recognizers in Rust architecture plugins

7 files changed, 141 insertions, 2 deletions

diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs
index 2c9843d9..edc16209 100644
--- a/rust/src/architecture.rs
+++ b/rust/src/architecture.rs
@@ -175,6 +175,7 @@ impl InstructionInfo {
     }
 }
 
+use crate::functionrecognizer::FunctionRecognizer;
 use crate::relocation::{CustomRelocationHandlerHandle, RelocationHandler};
 pub use binaryninjacore_sys::BNFlagRole as FlagRole;
 pub use binaryninjacore_sys::BNImplicitRegisterExtend as ImplicitRegisterExtend;
@@ -1646,6 +1647,13 @@ pub trait ArchitectureExt: Architecture {
     {
         crate::relocation::register_relocation_handler(self.as_ref(), name, func);
     }
+
+    fn register_function_recognizer<R>(&self, recognizer: R)
+    where
+        R: 'static + FunctionRecognizer + Send + Sync + Sized,
+    {
+        crate::functionrecognizer::register_arch_function_recognizer(self.as_ref(), recognizer);
+    }
 }
 
 impl<T: Architecture> ArchitectureExt for T {}
diff --git a/rust/src/function.rs b/rust/src/function.rs
index 0a01f337..7af9ee02 100644
--- a/rust/src/function.rs
+++ b/rust/src/function.rs
@@ -287,6 +287,20 @@ impl Function {
             Array::new(variables, count, ())
         }
     }
+
+    pub fn apply_imported_types(&self, sym: &Symbol, t: Option<&Type>) {
+        unsafe {
+            BNApplyImportedTypes(
+                self.handle,
+                sym.handle,
+                if let Some(t) = t {
+                    t.handle
+                } else {
+                    core::ptr::null_mut()
+                },
+            );
+        }
+    }
 }
 
 impl fmt::Debug for Function {
diff --git a/rust/src/functionrecognizer.rs b/rust/src/functionrecognizer.rs
new file mode 100644
index 00000000..ec567308
--- /dev/null
+++ b/rust/src/functionrecognizer.rs
@@ -0,0 +1,107 @@
+use crate::architecture::Architecture;
+use crate::{
+    architecture::CoreArchitecture, binaryview::BinaryView, function::Function, llil, mlil,
+};
+use binaryninjacore_sys::*;
+use std::os::raw::c_void;
+
+pub trait FunctionRecognizer {
+    fn recognize_low_level_il(
+        &self,
+        _bv: &BinaryView,
+        _func: &Function,
+        _llil: &llil::RegularFunction<CoreArchitecture>,
+    ) -> bool {
+        false
+    }
+
+    fn recognize_medium_level_il(
+        &self,
+        _bv: &BinaryView,
+        _func: &Function,
+        _mlil: &mlil::MediumLevelILFunction,
+    ) -> bool {
+        false
+    }
+}
+
+fn create_function_recognizer_registration<R>(recognizer: R) -> BNFunctionRecognizer
+where
+    R: 'static + FunctionRecognizer + Send + Sync + Sized,
+{
+    #[repr(C)]
+    struct FunctionRecognizerHandlerContext<R>
+    where
+        R: 'static + FunctionRecognizer + Send + Sync,
+    {
+        recognizer: R,
+    }
+
+    extern "C" fn cb_recognize_low_level_il<R>(
+        ctxt: *mut c_void,
+        bv: *mut BNBinaryView,
+        func: *mut BNFunction,
+        llil: *mut BNLowLevelILFunction,
+    ) -> bool
+    where
+        R: 'static + FunctionRecognizer + Send + Sync,
+    {
+        let custom_handler = unsafe { &*(ctxt as *mut R) };
+        let bv = unsafe { BinaryView::from_raw(BNNewViewReference(bv)) };
+        let arch = unsafe { BNGetFunctionArchitecture(func) };
+        let func = unsafe { Function::from_raw(BNNewFunctionReference(func)) };
+        if arch.is_null() {
+            return false;
+        }
+        let arch = unsafe { CoreArchitecture::from_raw(arch) };
+        let llil = unsafe { llil::RegularFunction::from_raw(arch, llil) };
+        custom_handler.recognize_low_level_il(bv.as_ref(), func.as_ref(), &llil)
+    }
+
+    extern "C" fn cb_recognize_medium_level_il<R>(
+        ctxt: *mut c_void,
+        bv: *mut BNBinaryView,
+        func: *mut BNFunction,
+        mlil: *mut BNMediumLevelILFunction,
+    ) -> bool
+    where
+        R: 'static + FunctionRecognizer + Send + Sync,
+    {
+        let custom_handler = unsafe { &*(ctxt as *mut R) };
+        let bv = unsafe { BinaryView::from_raw(BNNewViewReference(bv)) };
+        let func = unsafe { Function::from_raw(BNNewFunctionReference(func)) };
+        let mlil = unsafe { mlil::MediumLevelILFunction::from_raw(mlil) };
+        custom_handler.recognize_medium_level_il(bv.as_ref(), func.as_ref(), &mlil)
+    }
+
+    let recognizer = FunctionRecognizerHandlerContext { recognizer };
+    let raw = Box::into_raw(Box::new(recognizer));
+    BNFunctionRecognizer {
+        context: raw as *mut _,
+        recognizeLowLevelIL: Some(cb_recognize_low_level_il::<R>),
+        recognizeMediumLevelIL: Some(cb_recognize_medium_level_il::<R>),
+    }
+}
+
+pub fn register_global_function_recognizer<R>(recognizer: R)
+where
+    R: 'static + FunctionRecognizer + Send + Sync + Sized,
+{
+    let mut recognizer = create_function_recognizer_registration::<R>(recognizer);
+    unsafe {
+        BNRegisterGlobalFunctionRecognizer(&mut recognizer as *mut _);
+    }
+}
+
+pub(crate) fn register_arch_function_recognizer<R>(arch: &CoreArchitecture, recognizer: R)
+where
+    R: 'static + FunctionRecognizer + Send + Sync + Sized,
+{
+    let mut recognizer = create_function_recognizer_registration::<R>(recognizer);
+    unsafe {
+        BNRegisterArchitectureFunctionRecognizer(
+            arch.handle().as_ref().0,
+            &mut recognizer as *mut _,
+        );
+    }
+}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 3636eb40..5347491f 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -143,6 +143,7 @@ pub mod fileaccessor;
 pub mod filemetadata;
 pub mod flowgraph;
 pub mod function;
+pub mod functionrecognizer;
 pub mod headless;
 pub mod interaction;
 pub mod linearview;
diff --git a/rust/src/llil/instruction.rs b/rust/src/llil/instruction.rs
index c9817dbf..36733472 100644
--- a/rust/src/llil/instruction.rs
+++ b/rust/src/llil/instruction.rs
@@ -111,6 +111,7 @@ where
             LLIL_CALL | LLIL_CALL_STACK_ADJUST => {
                 InstrInfo::Call(Operation::new(self.function, op))
             }
+            LLIL_TAILCALL => InstrInfo::TailCall(Operation::new(self.function, op)),
             LLIL_SYSCALL => InstrInfo::Syscall(Operation::new(self.function, op)),
             LLIL_INTRINSIC => InstrInfo::Intrinsic(Operation::new(self.function, op)),
             _ => {
@@ -149,7 +150,7 @@ where
                 visit!(fb, &op.source_expr());
             }
             Push(ref op) => visit!(fb, &op.operand()),
-            Call(ref op) => visit!(fb, &op.target()),
+            Call(ref op) | TailCall(ref op) => visit!(fb, &op.target()),
             Intrinsic(ref _op) => {
                 // TODO: Use this when we support expression lists
                 // for expr in op.source_exprs() {
@@ -180,6 +181,7 @@ where
     JumpTo(Operation<'func, A, M, F, operation::JumpTo>),
 
     Call(Operation<'func, A, M, F, operation::Call>),
+    TailCall(Operation<'func, A, M, F, operation::Call>),
 
     Ret(Operation<'func, A, M, F, operation::Ret>),
     NoRet(Operation<'func, A, M, F, operation::NoArgs>),
diff --git a/rust/src/llil/mod.rs b/rust/src/llil/mod.rs
index f7696d57..8d199a2a 100644
--- a/rust/src/llil/mod.rs
+++ b/rust/src/llil/mod.rs
@@ -49,7 +49,7 @@ pub type LiftedExpr<'a, Arch> = Expression<'a, Arch, Mutable, NonSSA<LiftedNonSS
 pub type RegularFunction<Arch> = Function<Arch, Finalized, NonSSA<RegularNonSSA>>;
 pub type SSAFunction<Arch> = Function<Arch, Finalized, SSA>;
 
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, PartialEq, Eq)]
 pub enum Register<R: ArchReg> {
     ArchReg(R),
     Temp(u32),
diff --git a/rust/src/symbol.rs b/rust/src/symbol.rs
index 0e90c15c..bb9d157e 100644
--- a/rust/src/symbol.rs
+++ b/rust/src/symbol.rs
@@ -275,6 +275,13 @@ impl Symbol {
     pub fn external(&self) -> bool {
         self.binding() == Binding::Weak || self.binding() == Binding::Global
     }
+
+    pub fn imported_function_from_import_address_symbol(sym: &Symbol, addr: u64) -> Ref<Symbol> {
+        unsafe {
+            let res = BNImportedFunctionFromImportAddressSymbol(sym.handle, addr);
+            Symbol::ref_from_raw(res)
+        }
+    }
 }
 
 unsafe impl Send for Symbol {}