[WARP] Filter out LLIL_JUMP artifacts in lifted IL

With the new instruction retrieval we are getting the instructions at the end of the lifted function which are not really apart of the function, but they share the address of the last ret/jump.

2 files changed, 31 insertions, 9 deletions

diff --git a/plugins/warp/src/lib.rs b/plugins/warp/src/lib.rs
index 2172ac5d..2d512831 100644
--- a/plugins/warp/src/lib.rs
+++ b/plugins/warp/src/lib.rs
@@ -3,7 +3,7 @@ use crate::convert::{bn_comment_to_comment, bn_var_to_location, from_bn_symbol,
 use binaryninja::architecture::{
     Architecture, ImplicitRegisterExtend, Register as BNRegister, RegisterInfo,
 };
-use binaryninja::basic_block::BasicBlock as BNBasicBlock;
+use binaryninja::basic_block::{BasicBlock as BNBasicBlock};
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
 use binaryninja::confidence::MAX_CONFIDENCE;
 use binaryninja::function::{Function as BNFunction, NativeBlock};
@@ -186,7 +186,7 @@ pub fn basic_block_guid<M: FunctionMutability>(
             instr_bytes.truncate(instr_info.length);
 
             // Find variant and blacklisted instructions using lifted il.
-            for lifted_il_instr in lifted_il.instructions_at(instr_addr) {
+            for lifted_il_instr in filtered_instructions_at(lifted_il, instr_addr) {
                 // If instruction is blacklisted, don't include the bytes.
                 if is_blacklisted_instruction(&lifted_il_instr) {
                     continue;
@@ -209,7 +209,7 @@ pub fn basic_block_guid<M: FunctionMutability>(
             // TODO: A "mapped llil" or having some simple data flow, the simple data flow is the most attractive
             // TODO: "solution", but it would require
             if let Ok(llil) = &low_level_il {
-                for low_level_instr in llil.instructions_at(instr_addr) {
+                for low_level_instr in filtered_instructions_at(llil, instr_addr) {
                     if is_computed_variant_instruction(relocatable_regions, &low_level_instr) {
                         // Found a computed variant instruction, mask off the entire instruction.
                         instr_bytes.fill(0);
@@ -226,6 +226,25 @@ pub fn basic_block_guid<M: FunctionMutability>(
     BasicBlockGUID::from(basic_block_bytes.as_slice())
 }
 
+pub fn filtered_instructions_at<M: FunctionMutability>(
+    il: &LowLevelILFunction<M, NonSSA>,
+    addr: u64,
+) -> Vec<LowLevelILInstruction<M, NonSSA>> {
+    il.instructions_at(addr)
+        .into_iter()
+        .enumerate()
+        .take_while(|(i, instr)| match instr.kind() {
+            // Stop collecting instructions after we see a LLIL_RET, LLIL_NO_RET.
+            LowLevelILInstructionKind::NoRet(_) | LowLevelILInstructionKind::Ret(_) => false,
+            // Stop collecting instruction if we are probably the end function jump in lifted IL. This
+            // is emitted at the end of the function and will mess with our GUID.
+            LowLevelILInstructionKind::Jump(_) => *i != 0,
+            _ => true,
+        })
+        .map(|(_, instr)| instr)
+        .collect()
+}
+
 /// Is the instruction not included in the masked byte sequence?
 ///
 /// Blacklisted instructions will make an otherwise identical function GUID fail to match.
@@ -301,7 +320,7 @@ pub fn is_computed_variant_instruction<M: FunctionMutability>(
     instr: &LowLevelILInstruction<M, NonSSA>,
 ) -> bool {
     let is_expr_constant = |expr: &LowLevelILExpression<M, NonSSA, ValueExpr>| match expr.kind() {
-        LowLevelILExpressionKind::Const(_) => true,
+        LowLevelILExpressionKind::Const(_) | LowLevelILExpressionKind::ConstPtr(_) => true,
         _ => false,
     };
 
diff --git a/plugins/warp/src/plugin/render_layer.rs b/plugins/warp/src/plugin/render_layer.rs
index ab09093f..6c446e6f 100644
--- a/plugins/warp/src/plugin/render_layer.rs
+++ b/plugins/warp/src/plugin/render_layer.rs
@@ -1,6 +1,6 @@
 use crate::{
-    is_blacklisted_instruction, is_computed_variant_instruction, is_variant_instruction,
-    relocatable_regions,
+    filtered_instructions_at, is_blacklisted_instruction, is_computed_variant_instruction,
+    is_variant_instruction, relocatable_regions,
 };
 use binaryninja::basic_block::BasicBlock;
 use binaryninja::disassembly::DisassemblyTextLine;
@@ -51,7 +51,7 @@ impl HighlightRenderLayer {
         let relocatable_regions = relocatable_regions(&lifted_il.function().view());
         for line in lines {
             // We use address here instead of index since it's more reliable for other IL's.
-            for lifted_il_instr in lifted_il.instructions_at(line.address) {
+            for lifted_il_instr in filtered_instructions_at(lifted_il, line.address) {
                 if is_blacklisted_instruction(&lifted_il_instr) {
                     line.highlight = self.blacklist;
                 } else if is_variant_instruction(&relocatable_regions, &lifted_il_instr) {
@@ -59,7 +59,7 @@ impl HighlightRenderLayer {
                 }
             }
 
-            for llil_instr in llil.instructions_at(line.address) {
+            for llil_instr in filtered_instructions_at(llil, line.address) {
                 if is_computed_variant_instruction(&relocatable_regions, &llil_instr) {
                     line.highlight = self.computed_variant;
                 }
@@ -86,7 +86,10 @@ impl RenderLayer for HighlightRenderLayer {
     ) -> Vec<DisassemblyTextLine> {
         // Highlight any instruction that WARP will mask.
         let function = block.function();
-        if let (Ok(lifted_il), Ok(llil)) = (function.lifted_il(), function.low_level_il()) {
+        if let (Some(lifted_il), Some(llil)) = (
+            function.lifted_il_if_available(),
+            function.low_level_il_if_available(),
+        ) {
             self.highlight_lines(&lifted_il, &llil, &mut lines);
         }
         lines