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

2 files changed, 6 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;
                 }