Rust bindings for performing function lifting

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,83 @@ 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 inline_remapping: HashMap<ArchAndAddr, ArchAndAddr>>,
+    //pub user_indirect_branches: HashMap<ArchAndAddr, HashSet<ArchAndAddr>>>,
+    //pub auto_indirect_branches: HashMap<ArchAndAddr, HashSet<ArchAndAddr>>>,
+    //pub inlined_calls: HashSet<u64>,
 }
 
 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!(!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(
+                    *flc_ref.basicBlocks.add(i),
+                    NativeBlock::new(),
+                ))
+            };
+
+            blocks.push(block.unwrap());
+        }
+
+        let raw_no_return_calls: &[BNArchitectureAndAddress] =
+            std::slice::from_raw_parts(flc_ref.noReturnCalls, flc_ref.noReturnCallsCount);
+        let no_return_calls: HashSet<Location> =
+            raw_no_return_calls.iter().map(Location::from).collect();
 
-        FunctionLifterContext { handle }
+        let raw_contextual_return_locs: &[BNArchitectureAndAddress] = unsafe {
+            std::slice::from_raw_parts(
+                flc_ref.contextualFunctionReturnLocations,
+                flc_ref.contextualFunctionReturnCount,
+            )
+        };
+        let raw_contextual_return_vals: &[bool] = unsafe {
+            std::slice::from_raw_parts(
+                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,
+            function,
+            platform,
+            logger,
+            blocks,
+            no_return_calls,
+            contextual_returns,
+        }
+    }
+
+    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>(
@@ -1622,12 +1695,15 @@ where
         A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync,
     {
         let custom_arch = unsafe { &*(ctxt as *mut A) };
-        let function = unsafe {
-            LowLevelILMutableFunction::from_raw_with_arch(function, Some(*custom_arch.as_ref()))
+        let llil = unsafe {
+            LowLevelILMutableFunction::from_raw_with_arch(
+                BNNewLowLevelILFunctionReference(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

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/lifting.rs

@@ -23,7 +23,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 +1513,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;