[WARP] Fix possible skipped instructions when multiple IL expressions are appended for a given instruction

3 files changed, 31 insertions, 4 deletions

diff --git a/plugins/warp/src/lib.rs b/plugins/warp/src/lib.rs
index aab26932..852d64d6 100644
--- a/plugins/warp/src/lib.rs
+++ b/plugins/warp/src/lib.rs
@@ -155,7 +155,7 @@ pub fn basic_block_guid<M: FunctionMutability>(
             instr_bytes.truncate(instr_info.length);
 
             // Find variant and blacklisted instructions using lifted il.
-            if let Some(lifted_il_instr) = lifted_il.instruction_at(instr_addr) {
+            for lifted_il_instr in lifted_il.instructions_at(instr_addr) {
                 // If instruction is blacklisted, don't include the bytes.
                 if is_blacklisted_instruction(&lifted_il_instr) {
                     continue;
@@ -164,6 +164,7 @@ pub fn basic_block_guid<M: FunctionMutability>(
                 if is_variant_instruction(relocatable_regions, &lifted_il_instr) {
                     // Found a variant instruction, mask off the entire instruction.
                     instr_bytes.fill(0);
+                    break;
                 }
             }
 
@@ -177,10 +178,11 @@ 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 {
-                if let Some(low_level_instr) = llil.instruction_at(instr_addr) {
+                for low_level_instr in llil.instructions_at(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);
+                        break;
                     }
                 }
             }
diff --git a/plugins/warp/src/plugin/render_layer.rs b/plugins/warp/src/plugin/render_layer.rs
index d7984303..ab09093f 100644
--- a/plugins/warp/src/plugin/render_layer.rs
+++ b/plugins/warp/src/plugin/render_layer.rs
@@ -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.
-            if let Some(lifted_il_instr) = lifted_il.instruction_at(line.address) {
+            for lifted_il_instr in lifted_il.instructions_at(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 {
                 }
             }
 
-            if let Some(llil_instr) = llil.instruction_at(line.address) {
+            for llil_instr in llil.instructions_at(line.address) {
                 if is_computed_variant_instruction(&relocatable_regions, &llil_instr) {
                     line.highlight = self.computed_variant;
                 }
diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index f5fcf25f..cf3a432a 100644
--- a/rust/src/low_level_il/function.rs
+++ b/rust/src/low_level_il/function.rs
@@ -93,6 +93,31 @@ where
         }
     }
 
+    /// Get all the contiguous instructions for a given location.
+    ///
+    /// NOTE: This won't get you every instruction for a location, only the instructions
+    /// that are sequential from the starting instruction.
+    pub fn instructions_at<L: Into<Location>>(&self, loc: L) -> Vec<LowLevelILInstruction<M, F>> {
+        let loc = loc.into();
+        // TODO: Instructions sharing the same address are not always sequential.
+        // Gather all of the sequential instructions with the same address and same block.
+        self.instruction_index_at(loc)
+            .map(|mut idx| {
+                let mut instructions = Vec::new();
+                let block = self.basic_block_containing_index(idx);
+                while idx.0 < self.instruction_count() {
+                    let instr = LowLevelILInstruction::new(self, idx);
+                    if instr.address() != loc.addr || instr.basic_block() != block {
+                        break;
+                    }
+                    instructions.push(instr);
+                    idx = idx.next();
+                }
+                instructions
+            })
+            .unwrap_or_default()
+    }
+
     pub fn instruction_at<L: Into<Location>>(&self, loc: L) -> Option<LowLevelILInstruction<M, F>> {
         Some(LowLevelILInstruction::new(
             self,