Add support for Rust architecture hook plugins

arch/msp430/src/architecture.rs

@@ -75,7 +75,7 @@ impl Architecture for Msp430 {
         self.max_instr_len()
     }
 
-    fn associated_arch_by_addr(&self, _addr: u64) -> CoreArchitecture {
+    fn associated_arch_by_addr(&self, _addr: &mut u64) -> CoreArchitecture {
         self.handle
     }
 

arch/riscv/src/lib.rs

@@ -770,7 +770,7 @@ impl<D: RiscVDisassembler> Architecture for RiscVArch<D> {
         self.max_instr_len()
     }
 
-    fn associated_arch_by_addr(&self, _addr: u64) -> CoreArchitecture {
+    fn associated_arch_by_addr(&self, _addr: &mut u64) -> CoreArchitecture {
         self.handle
     }
 

rust/src/architecture.rs

@@ -16,7 +16,8 @@
 //! analysis and decompilation.
 //!
 //! For more information see the [`Architecture`] trait and the [`CoreArchitecture`] structure for
-//! querying already registered architectures.
+//! querying already registered architectures. To override selected behavior of an existing
+//! architecture, implement [`ArchitectureHook`] and call [`register_architecture_hook`].
 
 // RegisterInfo purge
 use binaryninjacore_sys::*;
@@ -52,7 +53,7 @@ use crate::relocation::{CustomRelocationHandlerHandle, RelocationHandler};
 use crate::basic_block::BasicBlock;
 use crate::confidence::Conf;
 use crate::logger::Logger;
-use crate::low_level_il::expression::ValueExpr;
+use crate::low_level_il::expression::{LowLevelExpressionIndex, ValueExpr};
 use crate::low_level_il::lifting::{
     get_default_flag_cond_llil, get_default_flag_write_llil, LowLevelILFlagWriteOp,
 };
@@ -61,6 +62,7 @@ use crate::low_level_il::{LowLevelILMutableExpression, LowLevelILMutableFunction
 pub mod basic_block;
 pub mod branches;
 pub mod flag;
+pub mod hook;
 pub mod instruction;
 pub mod intrinsic;
 pub mod register;
@@ -71,6 +73,7 @@ pub use basic_block::*;
 pub use binaryninjacore_sys::BNLinearSweepAnalysisCapability as LinearSweepAnalysisCapability;
 pub use branches::*;
 pub use flag::*;
+pub use hook::{register_architecture_hook, ArchitectureHook};
 pub use instruction::*;
 pub use intrinsic::*;
 pub use register::*;
@@ -180,7 +183,9 @@ pub trait Architecture: 'static + Sized + AsRef<CoreArchitecture> {
     /// In binaries with multiple architectures, you may wish to associate a specific architecture
     /// with a given virtual address. This can be seen in armv7 where odd addresses are associated
     /// with the thumb architecture.
-    fn associated_arch_by_addr(&self, _addr: u64) -> CoreArchitecture {
+    ///
+    /// The address can be updated in place to normalize it, such as clearing the ARM Thumb bit.
+    fn associated_arch_by_addr(&self, _addr: &mut u64) -> CoreArchitecture {
         *self.as_ref()
     }
 
@@ -234,11 +239,11 @@ pub trait Architecture: 'static + Sized + AsRef<CoreArchitecture> {
         self.instruction_text(data, addr)
     }
 
-    // TODO: Why do we need to return a boolean here? Does `None` not represent the same thing?
     /// Appends arbitrary low-level il instructions to `il`.
     ///
-    /// If `None` is returned, no instructions were appended and the data is invalid. If `Some` is returned,
-    /// the instructions consumed length is returned (necessary for variable length instruction decoding).
+    /// Returns the decoded length and whether lifting succeeded. `Some((length, false))` preserves
+    /// the decoded length even when lifting fails. `None` reports failure without updating the
+    /// caller's length.
     fn instruction_llil(
         &self,
         data: &[u8],
@@ -502,8 +507,12 @@ pub trait Architecture: 'static + Sized + AsRef<CoreArchitecture> {
         Vec::new()
     }
 
-    fn intrinsic_class(&self, _id: IntrinsicId) -> BNIntrinsicClass {
-        BNIntrinsicClass::GeneralIntrinsicClass
+    /// Returns the class from [`Intrinsic::class`], or the general class for an unknown ID.
+    fn intrinsic_class(&self, id: IntrinsicId) -> BNIntrinsicClass {
+        match self.intrinsic_from_id(id) {
+            Some(intrinsic) => intrinsic.class(),
+            None => BNIntrinsicClass::GeneralIntrinsicClass,
+        }
     }
 
     /// Get the [`Self::Intrinsic`] associated with the given [`IntrinsicId`].
@@ -893,8 +902,8 @@ impl Architecture for CoreArchitecture {
         unsafe { BNGetArchitectureOpcodeDisplayLength(self.handle) }
     }
 
-    fn associated_arch_by_addr(&self, addr: u64) -> CoreArchitecture {
-        let handle = unsafe { BNGetAssociatedArchitectureByAddress(self.handle, addr as *mut _) };
+    fn associated_arch_by_addr(&self, addr: &mut u64) -> CoreArchitecture {
+        let handle = unsafe { BNGetAssociatedArchitectureByAddress(self.handle, addr) };
         CoreArchitecture { handle }
     }
 
@@ -987,11 +996,7 @@ impl Architecture for CoreArchitecture {
             )
         };
 
-        if !success {
-            None
-        } else {
-            Some((size, true))
-        }
+        Some((size, success))
     }
 
     /// Performs basic block recovery and commits the results to the function analysis.
@@ -1020,12 +1025,29 @@ impl Architecture for CoreArchitecture {
 
     fn flag_write_llil<'a>(
         &self,
-        _flag: Self::Flag,
-        _flag_write: Self::FlagWrite,
-        _op: LowLevelILFlagWriteOp<Self::Register>,
-        _il: &'a LowLevelILMutableFunction,
+        flag: Self::Flag,
+        flag_write: Self::FlagWrite,
+        op: LowLevelILFlagWriteOp<Self::Register>,
+        il: &'a LowLevelILMutableFunction,
     ) -> Option<LowLevelILMutableExpression<'a, ValueExpr>> {
-        None
+        let (size, operation) = op.size_and_op();
+        let mut operands = op.raw_operands();
+        let index = unsafe {
+            BNGetArchitectureFlagWriteLowLevelIL(
+                self.handle,
+                operation,
+                size,
+                flag_write.id().0,
+                flag.id().0,
+                operands.as_mut_ptr(),
+                operands.len(),
+                il.handle,
+            )
+        };
+        Some(LowLevelILMutableExpression::new(
+            il,
+            LowLevelExpressionIndex(index),
+        ))
     }
 
     fn flags_required_for_flag_condition(
@@ -1058,19 +1080,36 @@ impl Architecture for CoreArchitecture {
 
     fn flag_cond_llil<'a>(
         &self,
-        _cond: FlagCondition,
-        _class: Option<Self::FlagClass>,
-        _il: &'a LowLevelILMutableFunction,
+        cond: FlagCondition,
+        class: Option<Self::FlagClass>,
+        il: &'a LowLevelILMutableFunction,
     ) -> Option<LowLevelILMutableExpression<'a, ValueExpr>> {
-        None
+        let index = unsafe {
+            BNGetArchitectureFlagConditionLowLevelIL(
+                self.handle,
+                cond,
+                class.map_or(0, |class| class.id().0),
+                il.handle,
+            )
+        };
+        Some(LowLevelILMutableExpression::new(
+            il,
+            LowLevelExpressionIndex(index),
+        ))
     }
 
     fn flag_group_llil<'a>(
         &self,
-        _group: Self::FlagGroup,
-        _il: &'a LowLevelILMutableFunction,
+        group: Self::FlagGroup,
+        il: &'a LowLevelILMutableFunction,
     ) -> Option<LowLevelILMutableExpression<'a, ValueExpr>> {
-        None
+        let index = unsafe {
+            BNGetArchitectureSemanticFlagGroupLowLevelIL(self.handle, group.id().0, il.handle)
+        };
+        Some(LowLevelILMutableExpression::new(
+            il,
+            LowLevelExpressionIndex(index),
+        ))
     }
 
     fn registers_all(&self) -> Vec<CoreRegister> {
@@ -1281,6 +1320,10 @@ impl Architecture for CoreArchitecture {
         }
     }
 
+    fn intrinsic_class(&self, id: IntrinsicId) -> BNIntrinsicClass {
+        unsafe { BNGetArchitectureIntrinsicClass(self.handle, id.0) }
+    }
+
     fn intrinsic_from_id(&self, id: IntrinsicId) -> Option<CoreIntrinsic> {
         CoreIntrinsic::new(*self, id)
     }
@@ -1553,10 +1596,19 @@ where
     A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync + Sized,
     F: FnOnce(CustomArchitectureHandle<A>, CoreArchitecture) -> A,
 {
-    register_architecture_impl(name, func, |_| {})
+    register_architecture_impl(ArchitectureRegistration::New(name), func, |_| {})
+}
+
+enum ArchitectureRegistration<'a> {
+    New(&'a str),
+    Hook(CoreArchitecture),
 }
 
-fn register_architecture_impl<A, F, C>(name: &str, func: F, customize: C) -> &'static A
+fn register_architecture_impl<A, F, C>(
+    registration: ArchitectureRegistration<'_>,
+    func: F,
+    customize: C,
+) -> &'static A
 where
     A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync + Sized,
     F: FnOnce(CustomArchitectureHandle<A>, CoreArchitecture) -> A,
@@ -1662,7 +1714,7 @@ where
         A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync,
     {
         let custom_arch = unsafe { &*(ctxt as *mut A) };
-        let addr = unsafe { *(addr) };
+        let addr = unsafe { &mut *addr };
 
         custom_arch.associated_arch_by_addr(addr).handle
     }
@@ -2204,28 +2256,27 @@ where
         A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync,
     {
         let custom_arch = unsafe { &*(ctxt as *mut A) };
-        let flag_write = custom_arch.flag_write_from_id(FlagWriteId(flag_write));
-        let flag = custom_arch.flag_from_id(FlagId(flag));
-        let operands = unsafe { std::slice::from_raw_parts(operands_raw, operand_count) };
+        let flag_write_type = custom_arch.flag_write_from_id(FlagWriteId(flag_write));
+        let flag_value = custom_arch.flag_from_id(FlagId(flag));
+        let operands = if operand_count == 0 {
+            &[]
+        } else {
+            unsafe { std::slice::from_raw_parts(operands_raw, operand_count) }
+        };
         let lifter = unsafe {
             LowLevelILMutableFunction::from_raw_with_arch(il, Some(*custom_arch.as_ref()))
         };
 
-        if let (Some(flag_write), Some(flag)) = (flag_write, flag) {
-            if let Some(op) = LowLevelILFlagWriteOp::from_op(custom_arch, size, op, operands) {
-                if let Some(expr) = custom_arch.flag_write_llil(flag, flag_write, op, &lifter) {
-                    // TODO verify that returned expr is a bool value
-                    return expr.index.0;
-                }
-            } else {
-                tracing::warn!(
-                    "unable to unpack flag write op: {:?} with {} operands",
-                    op,
-                    operands.len()
-                );
+        if let (Some(flag_write_type), Some(flag_value)) = (flag_write_type, flag_value) {
+            let decoded_op = LowLevelILFlagWriteOp::from_op(custom_arch, size, op, operands);
+            if let Some(expr) =
+                custom_arch.flag_write_llil(flag_value, flag_write_type, decoded_op, &lifter)
+            {
+                // TODO verify that returned expr is a bool value
+                return expr.index.0;
             }
 
-            let role = flag.role(flag_write.class());
+            let role = flag_value.role(flag_write_type.class());
 
             unsafe {
                 BNGetDefaultArchitectureFlagWriteLowLevelIL(
@@ -2413,11 +2464,7 @@ where
         A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync,
     {
         let custom_arch = unsafe { &*(ctxt as *mut A) };
-        match custom_arch.intrinsic_from_id(IntrinsicId(intrinsic)) {
-            Some(intrinsic) => intrinsic.class(),
-            // TODO: Make this unreachable?
-            None => BNIntrinsicClass::GeneralIntrinsicClass,
-        }
+        custom_arch.intrinsic_class(IntrinsicId(intrinsic))
     }
 
     extern "C" fn cb_intrinsic_name<A>(ctxt: *mut c_void, intrinsic: u32) -> *mut c_char
@@ -2726,8 +2773,6 @@ where
         custom_arch.skip_and_return_value(data, addr, val)
     }
 
-    let name = name.to_cstr();
-
     let uninit_arch = ArchitectureBuilder {
         arch: MaybeUninit::zeroed(),
         func: Some(func),
@@ -2825,8 +2870,23 @@ where
     customize(&mut custom_arch);
 
     unsafe {
-        let res = BNRegisterArchitecture(name.as_ptr(), &mut custom_arch as *mut _);
-        assert!(!res.is_null());
+        match registration {
+            ArchitectureRegistration::Hook(base) => {
+                // Hook registration swaps the callbacks and returns the previous implementation.
+                // Unlike BNRegisterArchitecture, it does not invoke the init callback.
+                let previous = BNRegisterArchitectureHook(base.handle, &mut custom_arch);
+                assert!(!previous.is_null());
+                cb_init::<A, F>(raw.cast(), previous);
+                // Finalization queries our callbacks to rebuild the register/flag caches, so the
+                // Rust implementation must be fully initialized before this call.
+                BNFinalizeArchitectureHook(base.handle);
+            }
+            ArchitectureRegistration::New(name) => {
+                let name = name.to_cstr();
+                let res = BNRegisterArchitecture(name.as_ptr(), &mut custom_arch);
+                assert!(!res.is_null());
+            }
+        }
 
         (*raw).arch.assume_init_mut()
     }
@@ -2901,7 +2961,7 @@ where
         true
     }
 
-    register_architecture_impl(name, func, |custom_arch| {
+    register_architecture_impl(ArchitectureRegistration::New(name), func, |custom_arch| {
         custom_arch.freeFunctionArchContext = Some(cb_free_function_arch_context_typed::<A>);
         custom_arch.getInstructionTextWithContext =
             Some(cb_get_instruction_text_with_context_typed::<A>);