Allow upgrading a mutable LowLevelILFunction to a finalized LowLevelILFunction through LowLevelILFunction::finalized in the Rust API

Author: emesare

plugins/lift_check/src/lib.rs

@@ -2,7 +2,7 @@ use binaryninja::architecture::{CoreArchitecture, Register, RegisterInfo};
 use binaryninja::logger::Logger;
 use binaryninja::low_level_il::expression::{ExpressionHandler, LowLevelILExpression, ValueExpr};
 use binaryninja::low_level_il::expression::LowLevelILExpressionKind::*;
-use binaryninja::low_level_il::function::{Finalized, LiftedNonSSA, LowLevelILFunction, NonSSA, RegularNonSSA};
+use binaryninja::low_level_il::function::{FunctionMutability, LiftedNonSSA, NonSSA};
 use binaryninja::low_level_il::instruction::LowLevelILInstructionKind::*;
 use log::{error, info, LevelFilter};
 use binaryninja::low_level_il::instruction::{InstructionHandler, LowLevelILInstruction};
@@ -19,7 +19,7 @@ const LIFT_CHECK_ACTIVITY_CONFIG: &str = r#"{
     }
 }"#;
 
-fn check_expression(expr: &LowLevelILExpression<CoreArchitecture, Finalized, NonSSA<LiftedNonSSA>, ValueExpr>,
+fn check_expression<M: FunctionMutability>(expr: &LowLevelILExpression<CoreArchitecture, M, NonSSA<LiftedNonSSA>, ValueExpr>,
                     required_size: Option<usize>)
 {
     match expr.kind() {
@@ -325,7 +325,7 @@ fn check_expression(expr: &LowLevelILExpression<CoreArchitecture, Finalized, Non
     }
 }
 
-fn check_instruction(inst: &LowLevelILInstruction<CoreArchitecture, Finalized, NonSSA<LiftedNonSSA>>) {
+fn check_instruction<M: FunctionMutability>(inst: &LowLevelILInstruction<CoreArchitecture, M, NonSSA<LiftedNonSSA>>) {
     match inst.kind() {
         SetReg(op) => {
             let required_expr_size = op.size();

rust/examples/workflow.rs

@@ -25,23 +25,21 @@ fn example_activity(analysis_context: &AnalysisContext) {
     );
     // If we have llil available, replace that as well.
     if let Some(llil) = unsafe { analysis_context.llil_function() } {
-        for basic_block in &func.basic_blocks() {
+        for basic_block in &llil.basic_blocks() {
             for instr in basic_block.iter() {
-                if let Some(llil_instr) = llil.instruction_at(instr) {
-                    llil_instr.visit_tree(&mut |expr| {
-                        if let LowLevelILExpressionKind::Const(_op) = expr.kind() {
-                            // Replace all consts with 0x1337.
-                            println!("Replacing llil expression @ 0x{:x} : {}", instr, expr.index);
-                            unsafe {
-                                llil.replace_expression(expr.index, llil.const_int(4, 0x1337))
-                            };
-                        }
-                        VisitorAction::Descend
-                    });
-                }
+                instr.visit_tree(&mut |expr| {
+                    if let LowLevelILExpressionKind::Const(_op) = expr.kind() {
+                        // Replace all consts with 0x1337.
+                        println!("Replacing llil expression @ 0x{:x} : {}", instr.address(), expr.index);
+                        unsafe {
+                            llil.replace_expression(expr.index, llil.const_int(4, 0x1337))
+                        };
+                    }
+                    VisitorAction::Descend
+                });
             }
         }
-        analysis_context.set_lifted_il_function(&llil);
+        analysis_context.set_llil_function(&llil.finalized());
     }
 }
 

rust/src/low_level_il/function.rs

@@ -156,17 +156,8 @@ where
             Function::ref_from_raw(func)
         }
     }
-}
 
-// LLIL basic blocks are not available until the function object
-// is finalized, so ensure we can't try requesting basic blocks
-// during lifting
-impl<A, F> LowLevelILFunction<A, Finalized, F>
-where
-    A: Architecture,
-    F: FunctionForm,
-{
-    pub fn basic_blocks(&self) -> Array<BasicBlock<LowLevelILBlock<A, Finalized, F>>> {
+    pub fn basic_blocks(&self) -> Array<BasicBlock<LowLevelILBlock<A, M, F>>> {
         use binaryninjacore_sys::BNGetLowLevelILBasicBlockList;
 
         unsafe {
@@ -204,6 +195,15 @@ impl LowLevelILFunction<CoreArchitecture, Mutable, NonSSA<LiftedNonSSA>> {
 
         unsafe { BNGenerateLowLevelILSSAForm(self.handle) };
     }
+    
+    /// Finalizes a mutable [`LowLevelILFunction`].
+    pub fn finalized(&self) -> Ref<LowLevelILFunction<CoreArchitecture, Finalized, NonSSA<LiftedNonSSA>>> {
+        use binaryninjacore_sys::BNFinalizeLowLevelILFunction;
+        unsafe { BNFinalizeLowLevelILFunction(self.handle) };
+        // We are transferring ownership of the handle to the caller as a finalized function.
+        // SAFETY: We know that the handle is valid and that the caller is responsible for the new reference. 
+        unsafe { LowLevelILFunction::from_raw(self.arch_handle, self.handle) }.to_owned()
+    }
 }
 
 impl<A, M, F> ToOwned for LowLevelILFunction<A, M, F>

rust/src/workflow.rs

@@ -67,10 +67,12 @@ impl AnalysisContext {
 
     // TODO: This returns LiftedNonSSA because the lifting code was written before we could patch the IL
     // TODO: At some point we need to take the lifting code and make it available to regular IL.
-    /// [`LowLevelILFunction`] used to represent Low Level IL
+    /// [`LowLevelILFunction`] used to represent Low Level IL, this is returned as mutable for the 
+    /// user to modify it is required to finalize the function using [`LowLevelILFunction::finalized`] 
+    /// before setting the function using [`LowLevelILFunction::set_llil_function`].
     pub unsafe fn llil_function<V: NonSSAVariant>(
         &self,
-    ) -> Option<Ref<LowLevelILFunction<CoreArchitecture, Finalized, NonSSA<V>>>> {
+    ) -> Option<Ref<LowLevelILFunction<CoreArchitecture, Mutable, NonSSA<V>>>> {
         let result = unsafe { BNAnalysisContextGetLowLevelILFunction(self.handle.as_ptr()) };
         let arch = self.function().arch();
         unsafe {
@@ -83,6 +85,8 @@ impl AnalysisContext {
 
     // TODO: This returns LiftedNonSSA because the lifting code was written before we could patch the IL
     // TODO: At some point we need to take the lifting code and make it available to regular IL.
+    /// To retrieve the finalized state of the function returned from [`AnalysisContext::llil_function`], use
+    /// [`LowLevelILFunction::finalized`] which insures that the function is in the correct state.
     pub fn set_llil_function<V: NonSSAVariant>(
         &self,
         value: &LowLevelILFunction<CoreArchitecture, Finalized, NonSSA<V>>,