Rust APIs for custom function lifters

8 files changed, 274 insertions, 16 deletions

diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs
index 76e9ab91..3d081b9f 100644
--- a/rust/src/architecture.rs
+++ b/rust/src/architecture.rs
@@ -27,7 +27,7 @@ use crate::{
     data_buffer::DataBuffer,
     disassembly::InstructionTextToken,
     ffi::INVALID_REGISTER,
-    function::Function,
+    function::{Function, Location, NativeBlock},
     platform::Platform,
     rc::*,
     relocation::CoreRelocationHandler,
@@ -35,6 +35,7 @@ use crate::{
     types::{NameAndType, Type},
     Endianness,
 };
+use std::collections::{HashMap, HashSet};
 use std::ops::Deref;
 use std::{
     borrow::Borrow,
@@ -48,7 +49,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,
@@ -592,13 +595,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();
 
-        FunctionLifterContext { handle }
+        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()
+        };
+
+        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>(
@@ -616,6 +739,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);
 
@@ -1630,12 +1761,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
@@ -2630,7 +2761,6 @@ where
 
     unsafe {
         let res = BNRegisterArchitecture(name.as_ptr(), &mut custom_arch as *mut _);
-
         assert!(!res.is_null());
 
         (*raw).arch.assume_init_mut()
diff --git a/rust/src/binary_view.rs b/rust/src/binary_view.rs
index c455635a..cda76ef2 100644
--- a/rust/src/binary_view.rs
+++ b/rust/src/binary_view.rs
@@ -1808,17 +1808,22 @@ 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 {
@@ -1826,8 +1831,13 @@ impl BinaryView {
                 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;
diff --git a/rust/src/logger.rs b/rust/src/logger.rs
index 66e594fd..d2903a5a 100644
--- a/rust/src/logger.rs
+++ b/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.
     ///
diff --git a/rust/src/low_level_il.rs b/rust/src/low_level_il.rs
index 12a74106..13c649b2 100644
--- a/rust/src/low_level_il.rs
+++ b/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>;
diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index 32e1f921..a0db838b 100644
--- a/rust/src/low_level_il/function.rs
+++ b/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> {
diff --git a/rust/src/low_level_il/lifting.rs b/rust/src/low_level_il/lifting.rs
index d162b3f6..202c9ed6 100644
--- a/rust/src/low_level_il/lifting.rs
+++ b/rust/src/low_level_il/lifting.rs
@@ -14,8 +14,10 @@
 
 use std::marker::PhantomData;
 
-use binaryninjacore_sys::{BNAddLowLevelILLabelForAddress, BNLowLevelILOperation};
-use binaryninjacore_sys::{BNLowLevelILLabel, BNRegisterOrConstant};
+use binaryninjacore_sys::{
+    BNAddLowLevelILLabelForAddress, BNLowLevelILClearIndirectBranches, BNLowLevelILLabel,
+    BNLowLevelILOperation, BNRegisterOrConstant, BNSetLowLevelILExprAttributes,
+};
 
 use super::*;
 use crate::architecture::{Architecture, FlagWriteId, RegisterId};
@@ -23,7 +25,8 @@ use crate::architecture::{CoreRegister, Register as ArchReg};
 use crate::architecture::{
     Flag, FlagClass, FlagCondition, FlagGroup, FlagRole, FlagWrite, Intrinsic,
 };
-use crate::function::Location;
+use crate::basic_block::BasicBlock;
+use crate::function::{Location, NativeBlock};
 
 pub trait LiftableLowLevelIL<'func> {
     type Result: ExpressionResultType;
@@ -1512,6 +1515,13 @@ impl LowLevelILMutableFunction {
         }
     }
 
+    pub fn set_current_source_block(&self, source: &BasicBlock<NativeBlock>) {
+        use binaryninjacore_sys::BNLowLevelILSetCurrentSourceBlock;
+        unsafe {
+            BNLowLevelILSetCurrentSourceBlock(self.handle, source.handle);
+        }
+    }
+
     pub fn label_for_address<L: Into<Location>>(&self, loc: L) -> Option<LowLevelILLabel> {
         use binaryninjacore_sys::BNGetLowLevelILLabelForAddress;
 
@@ -1561,6 +1571,25 @@ impl LowLevelILMutableFunction {
         }
         *label = new_label;
     }
+
+    pub fn set_expr_attributes(
+        &self,
+        expr: LowLevelExpressionIndex,
+        value: &ILInstructionAttributeSet,
+    ) {
+        let mut result = 0u32;
+        for flag in value {
+            result |= flag.value();
+        }
+
+        unsafe {
+            BNSetLowLevelILExprAttributes(self.handle, expr.0, result);
+        }
+    }
+
+    pub fn clear_indirect_branches(&self) {
+        unsafe { BNLowLevelILClearIndirectBranches(self.handle) };
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
diff --git a/rust/src/platform.rs b/rust/src/platform.rs
index e70e1277..9de125a1 100644
--- a/rust/src/platform.rs
+++ b/rust/src/platform.rs
@@ -20,8 +20,8 @@ use crate::{
     rc::*,
     string::*,
     types::{
-        QualifiedNameAndType, TypeContainer, TypeLibrary, TypeParserError, TypeParserErrorSeverity,
-        TypeParserResult,
+        QualifiedName, QualifiedNameAndType, Type, TypeContainer, TypeLibrary, TypeParserError,
+        TypeParserErrorSeverity, TypeParserResult,
     },
 };
 use binaryninjacore_sys::*;
@@ -279,6 +279,22 @@ impl Platform {
         }
     }
 
+    pub fn function_by_name<T: Into<QualifiedName>>(
+        &self,
+        name: T,
+        exact_match: bool,
+    ) -> Option<Ref<Type>> {
+        let mut raw_name = QualifiedName::into_raw(name.into());
+        unsafe {
+            let type_handle = BNGetPlatformFunctionByName(self.handle, &mut raw_name, exact_match);
+            QualifiedName::free_raw(raw_name);
+            if type_handle.is_null() {
+                return None;
+            }
+            Some(Type::ref_from_raw(type_handle))
+        }
+    }
+
     // TODO: system_calls
     // TODO: add a helper function to define a system call (platform function with a specific type)
 
diff --git a/rust/tests/platform.rs b/rust/tests/platform.rs
index 36a04d7f..09022ec1 100644
--- a/rust/tests/platform.rs
+++ b/rust/tests/platform.rs
@@ -22,6 +22,17 @@ fn test_platform_types() {
 }
 
 #[test]
+fn test_platform_function_by_name() {
+    let _session = Session::new().expect("Failed to initialize session");
+    let platform = Platform::by_name("windows-x86_64").expect("windows-x86_64 exists");
+    let platform_functions = platform.functions();
+    assert_ne!(platform_functions.len(), 0);
+
+    let function = platform_functions.get(0);
+    assert!(platform.function_by_name(function.name, true).is_some());
+}
+
+#[test]
 fn test_platform_calling_conventions() {
     let _session = Session::new().expect("Failed to initialize session");
     let platform = Platform::by_name("windows-x86_64").expect("windows-x86_64 exists");