[Rust] WIP: Improve situation around LLIL rewrites

Still working on it, will also need to refactor MLIL and HLIL to have the appropriate expression rewrite stuff

Author: emesare

rust/examples/workflow.rs

@@ -1,6 +1,6 @@
 use binaryninja::binary_view::BinaryViewExt;
 use binaryninja::low_level_il::expression::{ExpressionHandler, LowLevelILExpressionKind};
-use binaryninja::low_level_il::instruction::InstructionHandler;
+use binaryninja::low_level_il::instruction::{InstructionHandler, LowLevelInstructionIndex};
 use binaryninja::low_level_il::VisitorAction;
 use binaryninja::workflow::{Activity, AnalysisContext, Workflow};
 
@@ -32,8 +32,15 @@ fn example_activity(analysis_context: &AnalysisContext) {
                         if let LowLevelILExpressionKind::Const(_op) = expr.kind() {
                             // Replace all consts with 0x1337.
                             println!("Replacing llil expression @ 0x{:x} : {}", instr, expr.index);
+                            // When rewriting, expressions make sure to set the new expressions address.
+                            let src_operand = expr.kind().source_operand();
+                            println!("Source operand: {:?}", src_operand);
                             unsafe {
-                                llil.replace_expression(expr.index, llil.const_int(4, 0x1337))
+                                llil.replace_expression(
+                                    expr.index,
+                                    llil.const_int(4, 0x1337)
+                                        .with_address(expr.kind().address()),
+                                )
                             };
                         }
                         VisitorAction::Descend
@@ -68,16 +75,21 @@ pub fn main() {
     // traverse all llil expressions and look for the constant 0x1337
     for func in &bv.functions() {
         if let Ok(llil) = func.low_level_il() {
+            let last_llil_instr =
+                llil.instruction_from_index(LowLevelInstructionIndex(llil.instruction_count() - 1));
+            let last_llil_instr_addr = last_llil_instr.unwrap().address();
             for block in &llil.basic_blocks() {
                 for instr in block.iter() {
                     instr.visit_tree(&mut |expr| {
                         if let LowLevelILExpressionKind::Const(value) = expr.kind() {
                             if value.value() == 0x1337 {
                                 println!(
-                                    "Found constant 0x1337 at instruction 0x{:x} in function {}",
+                                    "Found constant 0x1337 at instruction 0x{:x} in function 0x{:x}",
                                     instr.address(),
                                     func.start()
                                 );
+                                // When rewriting, expressions make sure to set the address.
+                                assert_ne!(expr.kind().address(), last_llil_instr_addr, "Replaced expression address is not set");
                             }
                         }
                         VisitorAction::Descend

rust/src/architecture.rs

@@ -80,6 +80,18 @@ macro_rules! new_id_type {
             }
         }
 
+        impl From<u64> for $name {
+            fn from(value: u64) -> Self {
+                Self(value as $inner_type)
+            }
+        }
+
+        impl From<$name> for u64 {
+            fn from(value: $name) -> Self {
+                value.0 as u64
+            }
+        }
+
         impl std::fmt::Display for $name {
             fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                 write!(f, "{}", self.0)
@@ -118,21 +130,37 @@ pub trait Architecture: 'static + Sized + AsRef<CoreArchitecture> {
     fn max_instr_len(&self) -> usize;
     fn opcode_display_len(&self) -> usize;
 
-    fn associated_arch_by_addr(&self, addr: u64) -> CoreArchitecture;
+    fn associated_arch_by_addr(&self, _addr: u64) -> CoreArchitecture {
+        self.as_ref().handle()
+    }
 
+    /// Returns the [`InstructionInfo`] at the given virtual address with `data`.
+    ///
+    /// The [`InstructionInfo`] object should always fill the proper length and branches if not, the
+    /// next instruction will likely be incorrect.
     fn instruction_info(&self, data: &[u8], addr: u64) -> Option<InstructionInfo>;
+
+    /// Returns a of list of [`InstructionTextToken`]'s representing the instruction at the
+    /// given virtual address with `data`.
     fn instruction_text(
         &self,
         data: &[u8],
         addr: u64,
     ) -> Option<(usize, Vec<InstructionTextToken>)>;
+
+    /// Appends arbitrary low level il instructions representing the semantics of the instruction at
+    /// the given virtual address with `data`.
     fn instruction_llil(
         &self,
         data: &[u8],
         addr: u64,
         il: &LowLevelILMutableFunction,
     ) -> Option<(usize, bool)>;
 
+    /// Performs basic block recovery and commits the results to the function analysis.
+    ///
+    /// NOTE: Only implement this method if function-level analysis is required. Otherwise, do not
+    /// implement to let default basic block analysis take place.
     fn analyze_basic_blocks(
         &self,
         function: &mut Function,

rust/src/architecture/basic_block.rs

@@ -119,6 +119,7 @@ impl BasicBlockAnalysisContext {
         }
     }
 
+    /// Adds a contextual function return location and its value to the current function.
     pub fn add_contextual_return(&mut self, loc: impl Into<Location>, value: bool) {
         let loc = loc.into();
         if !self.contextual_returns.contains_key(&loc) {
@@ -128,20 +129,24 @@ impl BasicBlockAnalysisContext {
         self.contextual_returns.insert(loc, value);
     }
 
+    /// Adds a direct code reference to the current function.
     pub fn add_direct_code_reference(&mut self, target: u64, source: impl Into<Location>) {
         self.direct_code_references
             .entry(target)
             .or_insert(source.into());
     }
 
+    /// Adds a direct no-return call location to the current function.
     pub fn add_direct_no_return_call(&mut self, loc: impl Into<Location>) {
         self.direct_no_return_calls.insert(loc.into());
     }
 
+    /// Adds an address to the set of halted disassembly addresses.
     pub fn add_halted_disassembly_address(&mut self, loc: impl Into<Location>) {
         self.halted_disassembly_addresses.insert(loc.into());
     }
 
+    /// Creates a new [`BasicBlock`] at the specified address for the given [`CoreArchitecture`].
     pub fn create_basic_block(
         &self,
         arch: CoreArchitecture,
@@ -157,12 +162,14 @@ impl BasicBlockAnalysisContext {
         unsafe { Some(BasicBlock::ref_from_raw(raw_block, NativeBlock::new())) }
     }
 
+    /// Adds a [`BasicBlock`] to the current function.
     pub fn add_basic_block(&self, block: Ref<BasicBlock<NativeBlock>>) {
         unsafe {
             BNAnalyzeBasicBlocksContextAddBasicBlockToFunction(self.handle, block.handle);
         }
     }
 
+    /// Adds a temporary outgoing reference to the specified function.
     pub fn add_temp_outgoing_reference(&self, target: &Function) {
         unsafe {
             BNAnalyzeBasicBlocksContextAddTempReference(self.handle, target.handle);
@@ -241,6 +248,7 @@ impl BasicBlockAnalysisContext {
         }
     }
 
+    /// Finalizes the function's basic block analysis.
     pub fn finalize(&mut self) {
         if !self.direct_code_references.is_empty() {
             self.update_direct_code_references();

rust/src/low_level_il.rs

@@ -12,9 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use binaryninjacore_sys::BNLowLevelILInstruction;
 use std::borrow::Cow;
 use std::fmt;
-
+use std::fmt::{Display, Formatter};
 // TODO : provide some way to forbid emitting register reads for certain registers
 // also writing for certain registers (e.g. zero register must prohibit il.set_reg and il.reg
 // (replace with nop or const(0) respectively)
@@ -157,6 +158,30 @@ impl From<LowLevelILTempRegister> for LowLevelILRegisterKind<CoreRegister> {
     }
 }
 
+impl From<CoreRegister> for LowLevelILRegisterKind<CoreRegister> {
+    fn from(reg: CoreRegister) -> Self {
+        LowLevelILRegisterKind::Arch(reg)
+    }
+}
+
+impl PartialEq<CoreRegister> for LowLevelILRegisterKind<CoreRegister> {
+    fn eq(&self, other: &CoreRegister) -> bool {
+        match *self {
+            LowLevelILRegisterKind::Arch(ref r) => r == other,
+            LowLevelILRegisterKind::Temp(_) => false,
+        }
+    }
+}
+
+impl PartialEq<LowLevelILTempRegister> for LowLevelILRegisterKind<CoreRegister> {
+    fn eq(&self, other: &LowLevelILTempRegister) -> bool {
+        match *self {
+            LowLevelILRegisterKind::Arch(_) => false,
+            LowLevelILRegisterKind::Temp(ref r) => r == other,
+        }
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 pub enum LowLevelILSSARegisterKind<R: ArchReg> {
     Full {
@@ -209,3 +234,52 @@ pub enum VisitorAction {
     Sibling,
     Halt,
 }
+
+#[repr(transparent)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct LowLevelILOperandIndex(pub u32);
+
+impl LowLevelILOperandIndex {
+    pub fn next(&self) -> Self {
+        Self(self.0 + 1)
+    }
+}
+
+impl TryFrom<u32> for LowLevelILOperandIndex {
+    type Error = ();
+
+    fn try_from(value: u32) -> Result<Self, Self::Error> {
+        match value {
+            7 => Err(()),
+            value => Ok(Self(value)),
+        }
+    }
+}
+
+impl Display for LowLevelILOperandIndex {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.write_fmt(format_args!("{}", self.0))
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct LowLevelILSourceLocation {
+    pub address: u64,
+    /// The referenced source operand.
+    pub source_operand: Option<LowLevelILOperandIndex>,
+}
+
+impl LowLevelILSourceLocation {
+    pub fn raw_source_operand(&self) -> u32 {
+        self.source_operand.unwrap_or(LowLevelILOperandIndex(7)).0
+    }
+}
+
+impl From<&BNLowLevelILInstruction> for LowLevelILSourceLocation {
+    fn from(value: &BNLowLevelILInstruction) -> Self {
+        Self {
+            address: value.address,
+            source_operand: value.sourceOperand.try_into().ok(),
+        }
+    }
+}

rust/src/low_level_il/expression.rs

@@ -28,16 +28,18 @@ use std::fmt;
 use std::fmt::{Debug, Display, Formatter};
 use std::marker::PhantomData;
 
+pub trait ExpressionResultType: 'static + Debug {}
+
 /// Used as a marker for an [`LowLevelILExpression`] that **can** produce a value.
 #[derive(Copy, Clone, Debug)]
 pub struct ValueExpr;
 
+impl ExpressionResultType for ValueExpr {}
+
 /// Used as a marker for an [`LowLevelILExpression`] that can **not** produce a value.
 #[derive(Copy, Clone, Debug)]
 pub struct VoidExpr;
 
-pub trait ExpressionResultType: 'static + Debug {}
-impl ExpressionResultType for ValueExpr {}
 impl ExpressionResultType for VoidExpr {}
 
 #[repr(transparent)]
@@ -574,6 +576,10 @@ where
         self.raw_struct().address
     }
 
+    pub fn source_operand(&self) -> u32 {
+        self.raw_struct().sourceOperand
+    }
+
     /// Determines if the expressions represent the same operation
     ///
     /// It does not examine the operands for equality.

rust/src/low_level_il/function.rs

@@ -159,6 +159,17 @@ where
         }
     }
 
+    pub fn instruction_from_expr_index(
+        &self,
+        index: LowLevelExpressionIndex,
+    ) -> Option<LowLevelILInstruction<M, F>> {
+        if index.0 >= self.expression_count() {
+            None
+        } else {
+            Some(LowLevelILInstruction::new_with_expr_index(self, index))
+        }
+    }
+
     pub fn expression_count(&self) -> usize {
         unsafe {
             use binaryninjacore_sys::BNGetLowLevelILExprCount;