Rust APIs for custom function lifters

Author: zznop

rust/src/architecture.rs

@@ -26,14 +26,15 @@ use crate::{
     calling_convention::CoreCallingConvention,
     data_buffer::DataBuffer,
     disassembly::InstructionTextToken,
-    function::Function,
+    function::{Function, Location, NativeBlock},
     platform::Platform,
     rc::*,
     relocation::CoreRelocationHandler,
     string::{IntoCStr, *},
     types::{NameAndType, Type},
     Endianness,
 };
+use std::collections::{HashMap, HashSet};
 use std::ops::Deref;
 use std::{
     borrow::Borrow,
@@ -47,7 +48,9 @@ use std::ptr::NonNull;
 use crate::function_recognizer::FunctionRecognizer;
 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::lifting::{
     get_default_flag_cond_llil, get_default_flag_write_llil, LowLevelILFlagWriteOp,
@@ -591,13 +594,133 @@ pub trait ArchitectureWithFunctionContext: Architecture {
 
 pub struct FunctionLifterContext {
     pub(crate) handle: *mut BNFunctionLifterContext,
+    pub function: *mut BNLowLevelILFunction,
+    pub platform: Ref<Platform>,
+    pub logger: Ref<Logger>,
+    pub blocks: Vec<Ref<BasicBlock<NativeBlock>>>,
+    pub no_return_calls: HashSet<Location>,
+    pub contextual_returns: HashMap<Location, bool>,
+    pub inlined_remapping: HashMap<Location, Location>,
+    pub user_indirect_branches: HashMap<Location, HashSet<Location>>,
+    pub auto_indirect_branches: HashMap<Location, HashSet<Location>>,
+    //pub inlined_calls: HashSet<u64>,
+}
+
+unsafe fn lifter_context_slice<'a, T>(ptr: *const T, len: usize) -> &'a [T] {
+    if len == 0 {
+        &[]
+    } else {
+        debug_assert!(!ptr.is_null());
+        unsafe { std::slice::from_raw_parts(ptr, len) }
+    }
 }
 
 impl FunctionLifterContext {
-    pub unsafe fn from_raw(handle: *mut BNFunctionLifterContext) -> Self {
+    pub unsafe fn from_raw(
+        function: *mut BNLowLevelILFunction,
+        handle: *mut BNFunctionLifterContext,
+    ) -> Self {
+        debug_assert!(!function.is_null());
         debug_assert!(!handle.is_null());
+        let flc_ref = &*handle;
+        let platform = unsafe { Platform::ref_from_raw(BNNewPlatformReference(flc_ref.platform)) };
+        let logger = unsafe { Logger::ref_from_raw(BNNewLoggerReference(flc_ref.logger)) };
+
+        let mut blocks = Vec::new();
+        for i in 0..flc_ref.basicBlockCount {
+            let block = unsafe {
+                Some(BasicBlock::ref_from_raw(
+                    BNNewBasicBlockReference(*flc_ref.basicBlocks.add(i)),
+                    NativeBlock::new(),
+                ))
+            };
+
+            blocks.push(block.unwrap());
+        }
+
+        let raw_no_return_calls: &[BNArchitectureAndAddress] =
+            lifter_context_slice(flc_ref.noReturnCalls, flc_ref.noReturnCallsCount);
+        let no_return_calls: HashSet<Location> =
+            raw_no_return_calls.iter().map(Location::from).collect();
+
+        let raw_contextual_return_locs: &[BNArchitectureAndAddress] = unsafe {
+            lifter_context_slice(
+                flc_ref.contextualFunctionReturnLocations,
+                flc_ref.contextualFunctionReturnCount,
+            )
+        };
+        let raw_contextual_return_vals: &[bool] = unsafe {
+            lifter_context_slice(
+                flc_ref.contextualFunctionReturnValues,
+                flc_ref.contextualFunctionReturnCount,
+            )
+        };
+        let contextual_returns: HashMap<Location, bool> = raw_contextual_return_locs
+            .iter()
+            .map(Location::from)
+            .zip(raw_contextual_return_vals.iter().copied())
+            .collect();
+
+        let inlined_remapping: HashMap<Location, Location> = {
+            let raw_inline_remap_locs: &[BNArchitectureAndAddress] = lifter_context_slice(
+                flc_ref.inlinedRemappingKeys,
+                flc_ref.inlinedRemappingEntryCount,
+            );
+
+            let raw_inline_remap_dests: &[BNArchitectureAndAddress] = lifter_context_slice(
+                flc_ref.inlinedRemappingValues,
+                flc_ref.inlinedRemappingEntryCount,
+            );
+
+            raw_inline_remap_locs
+                .iter()
+                .map(Location::from)
+                .zip(raw_inline_remap_dests.iter().map(Location::from))
+                .collect()
+        };
 
-        FunctionLifterContext { handle }
+        let mut user_indirect_branches: HashMap<Location, HashSet<Location>> = HashMap::new();
+        let mut auto_indirect_branches: HashMap<Location, HashSet<Location>> = HashMap::new();
+        for i in 0..flc_ref.indirectBranchesCount {
+            let entry = unsafe { *flc_ref.indirectBranches.add(i) };
+            let src = Location::new(
+                Some(CoreArchitecture::from_raw(entry.sourceArch)),
+                entry.sourceAddr,
+            );
+            let dest = Location::new(
+                Some(CoreArchitecture::from_raw(entry.destArch)),
+                entry.destAddr,
+            );
+            if entry.autoDefined {
+                auto_indirect_branches.entry(src).or_default().insert(dest);
+            } else {
+                user_indirect_branches.entry(src).or_default().insert(dest);
+            }
+        }
+
+        FunctionLifterContext {
+            handle,
+            function: BNNewLowLevelILFunctionReference(function),
+            platform,
+            logger,
+            blocks,
+            no_return_calls,
+            contextual_returns,
+            inlined_remapping,
+            user_indirect_branches,
+            auto_indirect_branches,
+        }
+    }
+
+    pub fn prepare_block_translation(
+        &self,
+        func: &LowLevelILMutableFunction,
+        arch: &CoreArchitecture,
+        address: u64,
+    ) {
+        unsafe {
+            BNPrepareBlockTranslation(func.handle, arch.handle, address);
+        }
     }
 
     pub fn get_function_arch_context<A: ArchitectureWithFunctionContext>(
@@ -615,6 +738,14 @@ impl FunctionLifterContext {
     }
 }
 
+impl Drop for FunctionLifterContext {
+    fn drop(&mut self) {
+        if !self.function.is_null() {
+            unsafe { BNFreeLowLevelILFunction(self.function) };
+        }
+    }
+}
+
 // TODO: WTF?!?!?!?
 pub struct CoreArchitectureList(*mut *mut BNArchitecture, usize);
 
@@ -1622,12 +1753,12 @@ where
         A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync,
     {
         let custom_arch = unsafe { &*(ctxt as *mut A) };
-        let function = unsafe {
+        let llil = unsafe {
             LowLevelILMutableFunction::from_raw_with_arch(function, Some(*custom_arch.as_ref()))
         };
-        let mut context: FunctionLifterContext =
-            unsafe { FunctionLifterContext::from_raw(context) };
-        custom_arch.lift_function(function, &mut context)
+
+        let mut ctx = unsafe { FunctionLifterContext::from_raw(function, context) };
+        custom_arch.lift_function(llil, &mut ctx)
     }
 
     extern "C" fn cb_reg_name<A>(ctxt: *mut c_void, reg: u32) -> *mut c_char
@@ -2622,7 +2753,6 @@ where
 
     unsafe {
         let res = BNRegisterArchitecture(name.as_ptr(), &mut custom_arch as *mut _);
-
         assert!(!res.is_null());
 
         (*raw).arch.assume_init_mut()

rust/src/binary_view.rs

@@ -1807,26 +1807,36 @@ impl BinaryView {
         address: u64,
         platform: &Platform,
     ) -> Option<Ref<Function>> {
-        self.add_auto_function_ext(address, platform, None)
+        self.add_auto_function_ext(address, platform, None, false)
     }
 
     /// Add an auto function at the given `address` with the `platform` and function type.
     ///
+    /// The `auto_discovered` flag is used to prevent or allow this created function to be deleted if
+    /// it is never used (the function has no xrefs), if you are confident that this is a valid function
+    /// set this to `false`.
+    ///
     /// NOTE: If the view's default platform is not set, this will set it to `platform`.
     pub fn add_auto_function_ext(
         &self,
         address: u64,
         platform: &Platform,
         func_type: Option<&Type>,
+        auto_discovered: bool,
     ) -> Option<Ref<Function>> {
         unsafe {
             let func_type = match func_type {
                 Some(func_type) => func_type.handle,
                 None => std::ptr::null_mut(),
             };
 
-            let handle =
-                BNAddFunctionForAnalysis(self.handle, platform.handle, address, true, func_type);
+            let handle = BNAddFunctionForAnalysis(
+                self.handle,
+                platform.handle,
+                address,
+                auto_discovered,
+                func_type,
+            );
 
             if handle.is_null() {
                 return None;

rust/src/function.rs

@@ -514,7 +514,8 @@ impl Function {
         let arch = arch.unwrap_or_else(|| self.arch());
         let mut result = 0;
         unsafe {
-            BNGetFunctionBlockSortHint(self.handle, arch.handle, addr, &mut result).then_some(result)
+            BNGetFunctionBlockSortHint(self.handle, arch.handle, addr, &mut result)
+                .then_some(result)
         }
     }
 

rust/src/logger.rs

@@ -57,6 +57,14 @@ impl Logger {
         Self::new_with_session(name, LOGGER_DEFAULT_SESSION_ID)
     }
 
+    pub fn ref_from_raw(handle: *mut BNLogger) -> Ref<Logger> {
+        unsafe {
+            Ref::new(Logger {
+                handle: NonNull::new(handle).unwrap(),
+            })
+        }
+    }
+
     /// Create a logger scoped with the specific [`SessionId`], hiding the logs when the session
     /// is not active in the UI.
     ///

rust/src/low_level_il.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::borrow::Cow;
+use std::collections::HashSet;
 use std::fmt;
 use std::fmt::{Debug, Display};
 // TODO : provide some way to forbid emitting register reads for certain registers
@@ -316,3 +317,42 @@ pub enum VisitorAction {
     Sibling,
     Halt,
 }
+
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
+pub enum ILInstructionAttribute {
+    ILAllowDeadStoreElimination,
+    ILPreventDeadStoreElimination,
+    MLILAssumePossibleUse,
+    MLILUnknownSize,
+    SrcInstructionUsesPointerAuth,
+    ILPreventAliasAnalysis,
+    ILIsCFGProtected,
+    MLILPossiblyUnusedIntermediate,
+    HLILFoldableExpr,
+    HLILInvertableCondition,
+    HLILEarlyReturnPossible,
+    HLILSwitchRecoveryPossible,
+    ILTransparentCopy,
+}
+
+impl ILInstructionAttribute {
+    pub fn value(&self) -> u32 {
+        match self {
+            Self::ILAllowDeadStoreElimination => 1,
+            Self::ILPreventDeadStoreElimination => 2,
+            Self::MLILAssumePossibleUse => 4,
+            Self::MLILUnknownSize => 8,
+            Self::SrcInstructionUsesPointerAuth => 16,
+            Self::ILPreventAliasAnalysis => 32,
+            Self::ILIsCFGProtected => 64,
+            Self::MLILPossiblyUnusedIntermediate => 128,
+            Self::HLILFoldableExpr => 256,
+            Self::HLILInvertableCondition => 512,
+            Self::HLILEarlyReturnPossible => 1024,
+            Self::HLILSwitchRecoveryPossible => 2048,
+            Self::ILTransparentCopy => 4096,
+        }
+    }
+}
+
+pub type ILInstructionAttributeSet = HashSet<ILInstructionAttribute>;

rust/src/low_level_il/function.rs

@@ -209,6 +209,20 @@ where
             Some(unsafe { BasicBlock::ref_from_raw(block, LowLevelILBlock { function: self }) })
         }
     }
+
+    pub fn set_indirect_branches(&self, branches: &[Location]) {
+        let mut bn_branches: Box<[BNArchitectureAndAddress]> = branches
+            .iter()
+            .map(|loc| BNArchitectureAndAddress {
+                address: loc.addr,
+                arch: loc.arch.unwrap_or_else(|| self.arch()).handle,
+            })
+            .collect();
+
+        unsafe {
+            BNLowLevelILSetIndirectBranches(self.handle, bn_branches.as_mut_ptr(), branches.len());
+        }
+    }
 }
 
 impl<M: FunctionMutability> LowLevelILFunction<M, NonSSA> {